From bcfbff4eaf1a3180ea584ca9dbc65f44d5baec83 Mon Sep 17 00:00:00 2001 From: Garv Ghai Date: Sun, 19 Jul 2026 14:00:59 +0000 Subject: [PATCH 1/9] Add Kimi-K2.7 text backbone (DeepSeek-V3) to mstar MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements Kimi-K2.7's text model (HF DeepseekV3ForCausalLM) as a new, self-contained mstar model package. Golden-verified against vLLM references on a reduced config and validated end-to-end through the real paged FlashInfer cache (prefill + decode). Reuses existing mstar abstractions; no shared code modified. Changes added (mstar/model/kimi_k2_7/): - config.py: KimiK2Config — MLA latent dims, fine-grained MoE grouping, deepseek_yarn RoPE + reduced() test config. - MLA attention (naive/materialized) + deepseek_yarn RoPE (GPT-J interleaved). - Fine-grained MoE: group-limited sigmoid noaux_tc router + ungated shared expert (reuses the fused-expert GEMM). - Decoder layer (dense + MoE), KimiLanguageModel + KimiForCausalLM. - weight_loader.py: HF loader (name remap + stacked per-expert gate/up/down rules; MLA loaded by name); INT4 dequant-on-load designed and deferred. - submodules.py: KimiLLMSubmodule (ARNodeSubmodule prefill/decode lifecycle, YARN position_ids, CUDA-graph configs); real get_submodule (meta -> to_empty -> load_weights); registry entry + configs/kimi_k2_7.yaml. - Tests: dummy-mode modular + GPU integration goldens (components, MoE, MLA, full forward, real paged attention, weight loading, submodule e2e). Major issues resolved: - FlashInfer SM90 prefill static_asserts head_dim_vo in {64,128,256}, so the naive-MLA head_dim (192 real / 24 reduced) won't JIT-build. Mitigation: pad q/k/v to the next supported dim and fold a compensating softmax scale (mscale^2 * sqrt(padded/qk_head_dim)); golden-verified vs SDPA. - KimiYarnRotaryEmbedding computed inv_freq as an __init__ buffer that meta->to_empty leaves uninitialized (would silently corrupt YARN on the real load path). Fix: compute lazily; audit confirms the loaded model has zero stray buffers. --- .gitignore | 5 + configs/kimi_k2_7.yaml | 12 + mstar/model/kimi_k2_7/__init__.py | 0 mstar/model/kimi_k2_7/components/__init__.py | 0 mstar/model/kimi_k2_7/components/attention.py | 152 +++++++ mstar/model/kimi_k2_7/components/causal_lm.py | 102 +++++ .../kimi_k2_7/components/decoder_layer.py | 73 ++++ .../kimi_k2_7/components/language_model.py | 116 ++++++ mstar/model/kimi_k2_7/components/moe.py | 258 ++++++++++++ mstar/model/kimi_k2_7/components/rope.py | 146 +++++++ mstar/model/kimi_k2_7/config.py | 151 +++++++ mstar/model/kimi_k2_7/kimi_model.py | 386 ++++++++++++++++++ mstar/model/kimi_k2_7/submodules.py | 269 ++++++++++++ mstar/model/kimi_k2_7/weight_loader.py | 204 +++++++++ mstar/model/registry.py | 5 + test/integration/test_kimi_components.py | 148 +++++++ test/integration/test_kimi_decoder_layer.py | 261 ++++++++++++ .../test_kimi_flashinfer_attention.py | 169 ++++++++ test/integration/test_kimi_forward.py | 271 ++++++++++++ test/integration/test_kimi_mla.py | 197 +++++++++ test/integration/test_kimi_mla_paged.py | 201 +++++++++ test/integration/test_kimi_moe.py | 213 ++++++++++ test/integration/test_kimi_submodule.py | 308 ++++++++++++++ test/integration/test_kimi_weight_loading.py | 235 +++++++++++ test/modular/test_kimi_model.py | 106 +++++ 25 files changed, 3988 insertions(+) create mode 100644 configs/kimi_k2_7.yaml create mode 100644 mstar/model/kimi_k2_7/__init__.py create mode 100644 mstar/model/kimi_k2_7/components/__init__.py create mode 100644 mstar/model/kimi_k2_7/components/attention.py create mode 100644 mstar/model/kimi_k2_7/components/causal_lm.py create mode 100644 mstar/model/kimi_k2_7/components/decoder_layer.py create mode 100644 mstar/model/kimi_k2_7/components/language_model.py create mode 100644 mstar/model/kimi_k2_7/components/moe.py create mode 100644 mstar/model/kimi_k2_7/components/rope.py create mode 100644 mstar/model/kimi_k2_7/config.py create mode 100644 mstar/model/kimi_k2_7/kimi_model.py create mode 100644 mstar/model/kimi_k2_7/submodules.py create mode 100644 mstar/model/kimi_k2_7/weight_loader.py create mode 100644 test/integration/test_kimi_components.py create mode 100644 test/integration/test_kimi_decoder_layer.py create mode 100644 test/integration/test_kimi_flashinfer_attention.py create mode 100644 test/integration/test_kimi_forward.py create mode 100644 test/integration/test_kimi_mla.py create mode 100644 test/integration/test_kimi_mla_paged.py create mode 100644 test/integration/test_kimi_moe.py create mode 100644 test/integration/test_kimi_submodule.py create mode 100644 test/integration/test_kimi_weight_loading.py create mode 100644 test/modular/test_kimi_model.py diff --git a/.gitignore b/.gitignore index 5b5cc3a84..16d9bea36 100644 --- a/.gitignore +++ b/.gitignore @@ -29,3 +29,8 @@ mstar/worker/ASYNC_REDESIGN.md # local AI-assistant context (kept local, not published — cf. vllm-omni) CLAUDE.md AGENTS.md +mstar_traces/ +.claude/skills/ + +# local golden-extraction harness (Kimi-K2.7 port dev tooling; not for main) +tools/kimi_goldens/ diff --git a/configs/kimi_k2_7.yaml b/configs/kimi_k2_7.yaml new file mode 100644 index 000000000..7291213b9 --- /dev/null +++ b/configs/kimi_k2_7.yaml @@ -0,0 +1,12 @@ +model: "kimi_k2_7" +# Kimi-K2.7 text backbone (DeepSeek-V3). This single-GPU config (tp_size 1) is the +# bring-up / reduced-model deployment. The real 1T-param MoE needs TP8 across a +# node — set `ranks: [0,1,2,3,4,5,6,7]` and `tp_size: 8` — and, being >1 GPU, the +# single-node tensor transport TENSOR_PROTOCOL=SHM (or TCP; never RDMA on +# coriander). One KV_CACHE LLM node with the prefill + decode(Loop) walks. +max_seq_len: 262144 +node_groups: + - node_names: [LLM] + ranks: [0] + tp_size: 1 + graph_walks: [prefill, decode] diff --git a/mstar/model/kimi_k2_7/__init__.py b/mstar/model/kimi_k2_7/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/mstar/model/kimi_k2_7/components/__init__.py b/mstar/model/kimi_k2_7/components/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/mstar/model/kimi_k2_7/components/attention.py b/mstar/model/kimi_k2_7/components/attention.py new file mode 100644 index 000000000..8141f9eb5 --- /dev/null +++ b/mstar/model/kimi_k2_7/components/attention.py @@ -0,0 +1,152 @@ +"""Kimi-K2.7 / DeepSeek-V3 MLA attention — naive / materialized path. + +MLA compresses q and k/v through low-rank latents, then (in the naive path) +projects the latent back up to full per-head K/V and runs ordinary attention. +This avoids the weight-absorbed path (``W_UK``/``W_UV``) and its bespoke kernel — +throughput caveat noted — so it drops straight onto mstar's paged +``run_attention`` ``[tokens, heads, head_dim]`` interface, matching vLLM's +``DeepseekV2Attention`` (the non-absorbed class). + +Per-token shape story (H heads, Dnope=qk_nope, Drope=qk_rope, Dqk=Dnope+Drope, +Dv=v_head_dim, L=kv_lora_rank): + - q: ``q_a_proj`` -> ``q_a_layernorm`` -> ``q_b_proj`` -> ``[T,H,Dqk]``, split + into ``q_nope[..,Dnope]`` / ``q_pe[..,Drope]``. + - kv: ``kv_a_proj_with_mqa`` -> ``[L | Drope]``; the ``L`` slice is RMS-normed + and ``kv_b_proj``-ed to per-head ``[k_nope[..,Dnope] | v[..,Dv]]``; the trailing + ``Drope`` slice is the single shared MQA rope key ``k_pe[T,1,Drope]``. + - YARN RoPE rotates only ``q_pe`` (per head) and ``k_pe`` (broadcast to H heads). + - assemble ``k = [k_nope | k_pe_broadcast] -> [T,H,Dqk]``; zero-pad ``q``/``k`` + (Dqk) and ``v`` (Dv) up to ``padded_head_dim`` (M6 mitigation — FlashInfer SM90 + rejects ``head_dim_vo`` not in {64,128,256}); fold the scale boost into ``q`` + (``run_attention`` uses the fixed ``1/sqrt(padded_head_dim)`` scale), attend, + slice the output back to ``Dv``, ``o_proj``. + +Cache config for this node: ``num_kv_heads == num_qo_heads == num_attention_heads``, +``head_dim == padded_head_dim`` (256 for the real Dqk=192, 64 for the reduced +Dqk=24). Weight-absorbed MLA (native latent dims, no pad) and the +``fused_qkv_a_proj`` weight fusion are deferred to the perf backlog. +""" +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F +from torch import nn + +from mstar.distributed.communication import CommGroup +from mstar.engine.cache_manager import BatchedCacheManager +from mstar.model.components.distributed import ColumnParallelLinear, RowParallelLinear +from mstar.model.components.norm import RMSNorm +from mstar.model.kimi_k2_7.components.rope import KimiYarnRotaryEmbedding, yarn_get_mscale +from mstar.model.kimi_k2_7.config import KimiK2Config + + +class KimiMLAAttention(nn.Module): + """Multi-head Latent Attention (naive/materialized).""" + + def __init__(self, config: KimiK2Config, comm_group: CommGroup | None = None) -> None: + super().__init__() + if comm_group is None: + comm_group = CommGroup.trivial() + + # TODO(M6): shard num_heads across TP; naive path assumes local == total. + self.num_heads = config.num_attention_heads + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.qk_head_dim = config.qk_head_dim + self.v_head_dim = config.v_head_dim + self.kv_lora_rank = config.kv_lora_rank + # FlashInfer SM90 rejects head_dim_vo not in {64,128,256}, so q/k/v are + # zero-padded to this width for the paged run_attention (M6 mitigation); + # the attention output is sliced back to v_head_dim. See config docstring. + self.padded_head_dim = config.padded_head_dim + h = self.num_heads + + # Q: two-stage low-rank (q_a down -> norm -> q_b up). Down-projections are + # replicated (small rank); up-projections shard over heads under TP. + self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=False) + self.q_a_layernorm = RMSNorm(config.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + comm_group, config.q_lora_rank, h * self.qk_head_dim, bias=False) + + # KV: shared latent + decoupled rope key. + self.kv_a_proj_with_mqa = nn.Linear( + config.hidden_size, config.kv_lora_rank + config.qk_rope_head_dim, bias=False) + self.kv_a_layernorm = RMSNorm(config.kv_lora_rank, eps=config.rms_norm_eps) + self.kv_b_proj = ColumnParallelLinear( + comm_group, config.kv_lora_rank, + h * (config.qk_nope_head_dim + config.v_head_dim), bias=False) + + self.o_proj = RowParallelLinear( + comm_group, h * config.v_head_dim, config.hidden_size, + bias=False, input_is_parallel=True, reduce_results=True) + + rope = config.rope_scaling + self.rotary = KimiYarnRotaryEmbedding( + rotary_dim=config.qk_rope_head_dim, + base=config.rope_theta, + factor=rope["factor"], + original_max_position_embeddings=rope["original_max_position_embeddings"], + beta_fast=rope.get("beta_fast", 32), + beta_slow=rope.get("beta_slow", 1), + mscale=rope.get("mscale", 1.0), + mscale_all_dim=rope.get("mscale_all_dim", 0.0), + ) + # Softmax-scale boost folded into q because run_attention applies a fixed + # 1/sqrt(head_dim) scale and exposes no custom sm_scale. DeepSeek's intended + # softmax scale is ``qk_head_dim**-0.5 * mscale**2``; run_attention now runs + # over the PADDED head dim, so it uses ``padded_head_dim**-0.5``. The + # zero-pad dims contribute 0 to q·k, so to recover the intended scale we + # fold ``mscale**2 * sqrt(padded_head_dim / qk_head_dim)`` into q: + # scores = (q*boost)·k * padded_head_dim**-0.5 + # = q·k * mscale**2 * sqrt(padded/qk) * padded**-0.5 + # = q·k * mscale**2 * qk**-0.5 (the DeepSeek scale). + mscale = yarn_get_mscale(rope["factor"], rope.get("mscale_all_dim", 0.0)) + self.softmax_scale_boost = ( + mscale * mscale * math.sqrt(self.padded_head_dim / self.qk_head_dim) + ) + + def forward( + self, + hidden_states: torch.Tensor, + cache_handle: BatchedCacheManager, + position_ids: torch.Tensor, + ) -> torch.Tensor: + num_tokens = hidden_states.shape[0] + h = self.num_heads + + # --- Q --- + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(num_tokens, h, self.qk_head_dim) + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + # --- KV latent --- + latent = self.kv_a_proj_with_mqa(hidden_states) # (T, L + Drope) + kv_a, k_pe = latent.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + kv = self.kv_b_proj(self.kv_a_layernorm(kv_a)) + kv = kv.view(num_tokens, h, self.qk_nope_head_dim + self.v_head_dim) + k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) + k_pe = k_pe.view(num_tokens, 1, self.qk_rope_head_dim) # shared MQA rope key + + # --- RoPE (only the pe slices) --- + q_pe, k_pe = self.rotary(position_ids, q_pe, k_pe) + + # --- assemble full q / k (k_pe broadcast over heads) --- + q = torch.cat([q_nope, q_pe], dim=-1) # (T, H, Dqk) + k_pe = k_pe.expand(num_tokens, h, self.qk_rope_head_dim) + k = torch.cat([k_nope, k_pe], dim=-1) # (T, H, Dqk) + + # --- zero-pad q/k (Dqk) and v (Dv) up to padded_head_dim for the paged + # run_attention (FlashInfer SM90 requires head_dim_vo in {64,128,256}) --- + qk_pad = self.padded_head_dim - self.qk_head_dim + q = F.pad(q, [0, qk_pad]) # (T, H, Dpad) + k = F.pad(k, [0, qk_pad]) # (T, H, Dpad) + v = F.pad(v, [0, self.padded_head_dim - self.v_head_dim]) # (T, H, Dpad) + + # --- softmax boost folded into q (compensates padded_head_dim scale), + # attend, strip the pad + v-pad, project --- + q = q * self.softmax_scale_boost + attn = cache_handle.run_attention(q=q, k=k, v=v) # (T, H, Dpad) + attn = attn[..., : self.v_head_dim].reshape(num_tokens, h * self.v_head_dim) + return self.o_proj(attn) diff --git a/mstar/model/kimi_k2_7/components/causal_lm.py b/mstar/model/kimi_k2_7/components/causal_lm.py new file mode 100644 index 000000000..7d45f517d --- /dev/null +++ b/mstar/model/kimi_k2_7/components/causal_lm.py @@ -0,0 +1,102 @@ +"""Kimi-K2.7 / DeepSeek-V3 assembled text backbone (M4 assembly). + +Stacks the M4 :class:`KimiDecoderLayer` blocks between a token embedding and a +final RMSNorm (:class:`KimiLanguageModel`), then wraps that with the untied LM +head (:class:`KimiForCausalLM`). This is the full text forward: token ids → +logits. + +The per-layer cache-handle contract mirrors ``OrpheusLanguageModel`` exactly: +each layer is preceded by ``cache_handle.set_layer_idx(layer_idx)`` (so the paged +KV cache writes/reads the right layer slice), and the loop is followed by a +single ``cache_handle.advance_seq_lens()`` (so every request's ``seq_len`` / +``position_id_start`` steps forward once per forward pass, not once per layer). +The naive MLA reads ``position_ids`` for its YARN RoPE, so unlike Orpheus we +thread ``position_ids`` through each layer. + +Lives in its own module (not ``language_model.py``) to keep the import graph +acyclic: ``decoder_layer`` imports the ``language_model`` builders, so the +assembly that imports ``decoder_layer`` must sit downstream of both. +""" +from __future__ import annotations + +import torch +from torch import nn + +from mstar.distributed.communication import CommGroup +from mstar.engine.cache_manager import BatchedCacheManager +from mstar.model.kimi_k2_7.components.decoder_layer import KimiDecoderLayer +from mstar.model.kimi_k2_7.components.language_model import ( + build_embedding, + build_lm_head, + build_rmsnorm, +) +from mstar.model.kimi_k2_7.config import KimiK2Config + + +class KimiLanguageModel(nn.Module): + """Embedding + stacked decoder layers + final norm (returns hidden states).""" + + def __init__( + self, config: KimiK2Config, comm_group: CommGroup | None = None + ) -> None: + super().__init__() + self.embed_tokens = build_embedding(config, comm_group=comm_group) + self.layers = nn.ModuleList( + [ + KimiDecoderLayer(config, layer_idx, comm_group=comm_group) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = build_rmsnorm(config) + + def forward( + self, + input_ids: torch.Tensor, + cache_handle: BatchedCacheManager, + position_ids: torch.Tensor, + ) -> torch.Tensor: + hidden_states = self.embed_tokens(input_ids) + for layer_idx, decoder_layer in enumerate(self.layers): + cache_handle.set_layer_idx(layer_idx) + hidden_states = decoder_layer( + hidden_states, cache_handle, position_ids + ) + cache_handle.advance_seq_lens() + return self.norm(hidden_states) + + +class KimiForCausalLM(nn.Module): + """Text backbone + untied LM head (returns ``[..., vocab]`` logits).""" + + def __init__( + self, config: KimiK2Config, comm_group: CommGroup | None = None + ) -> None: + super().__init__() + self.config = config + self.model = KimiLanguageModel(config, comm_group=comm_group) + self.lm_head = build_lm_head(config, comm_group=comm_group) + + def forward( + self, + input_ids: torch.Tensor, + cache_handle: BatchedCacheManager, + position_ids: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, cache_handle, position_ids) + return self.lm_head(hidden_states) + + def load_weights(self, weights, **kwargs) -> set[str]: + """Load an HF DeepSeek-V3 checkpoint stream (M5). + + Called by the shared ``mstar.model.loader.load_weights(model, source, + device)`` driver (mirrors ``OrpheusForCausalLM.load_weights``). Delegates + to :func:`mstar.model.kimi_k2_7.weight_loader.load_kimi_hf_weights` for + the Kimi remap + fused-expert stacked rules. Returns the set of loaded + param paths. + """ + from mstar.model.kimi_k2_7.weight_loader import load_kimi_hf_weights + + return load_kimi_hf_weights( + self, weights, self.config.n_routed_experts, + ) diff --git a/mstar/model/kimi_k2_7/components/decoder_layer.py b/mstar/model/kimi_k2_7/components/decoder_layer.py new file mode 100644 index 000000000..88e581d01 --- /dev/null +++ b/mstar/model/kimi_k2_7/components/decoder_layer.py @@ -0,0 +1,73 @@ +"""Kimi-K2.7 / DeepSeek-V3 decoder layer (M4 assembly). + +One pre-norm transformer block: MLA self-attention then a feed-forward that is +either the dense SwiGLU MLP (the ``first_k_dense_replace`` early layers) or the +fine-grained sigmoid-routed MoE block. Both feed-forwards expose the same +``(x) -> x`` interface, so the residual wiring here is agnostic to which it holds +(``build_mlp_for_layer`` picks per ``layer_idx``). + +This is a Kimi-specific decoder layer rather than the shared +``mstar.model.components.DecoderLayer`` because MLA attention needs +``position_ids`` threaded through its forward (the shared layer's +``self_attn(x, cache_handle=...)`` signature has no position channel — YARN RoPE +is applied inside the attention over the decoupled ``qk_rope`` slice). + +Residual structure mirrors vLLM ``DeepseekV2DecoderLayer.forward``: + residual = h + h = input_layernorm(h); h = self_attn(h, cache, pos); h = residual + h + residual = h + h = post_attention_layernorm(h); h = mlp(h); h = residual + h +""" +from __future__ import annotations + +import torch +from torch import nn + +from mstar.distributed.communication import CommGroup +from mstar.engine.cache_manager import BatchedCacheManager +from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention +from mstar.model.kimi_k2_7.components.language_model import ( + build_mlp_for_layer, + build_rmsnorm, +) +from mstar.model.kimi_k2_7.config import KimiK2Config + + +class KimiDecoderLayer(nn.Module): + """Pre-norm MLA + (dense-or-MoE) feed-forward block. + + Args: + config: model config. + layer_idx: index into the stack; selects the dense MLP (``layer_idx < + first_k_dense_replace``) or the MoE block (``build_mlp_for_layer``). + comm_group: TP comm group (trivial single-rank if ``None``). + """ + + def __init__( + self, + config: KimiK2Config, + layer_idx: int, + comm_group: CommGroup | None = None, + ) -> None: + super().__init__() + self.layer_idx = layer_idx + self.self_attn = KimiMLAAttention(config, comm_group=comm_group) + self.mlp = build_mlp_for_layer(config, layer_idx, comm_group=comm_group) + self.input_layernorm = build_rmsnorm(config) + self.post_attention_layernorm = build_rmsnorm(config) + + def forward( + self, + hidden_states: torch.Tensor, + cache_handle: BatchedCacheManager, + position_ids: torch.Tensor, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, cache_handle, position_ids) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states diff --git a/mstar/model/kimi_k2_7/components/language_model.py b/mstar/model/kimi_k2_7/components/language_model.py new file mode 100644 index 000000000..6edcf5ef1 --- /dev/null +++ b/mstar/model/kimi_k2_7/components/language_model.py @@ -0,0 +1,116 @@ +"""Kimi-K2.7 language-model components (DeepSeek-V3 text backbone). + +M1 (cheap reuse) wiring: builders that map ``KimiK2Config`` onto mstar's existing +reused primitives — token embedding, the dense SwiGLU MLP (the +``first_k_dense_replace`` early layers), RMSNorm, and the LM head. These are the +pieces DeepSeek-V3 shares verbatim with a standard Llama-style stack; the +Kimi-specific parts (MLA attention, fine-grained sigmoid-routed MoE, YARN RoPE) +are separate milestones (M2/M3). + +Each builder is thin on purpose: it fixes the config→component mapping (dims, +``silu`` activation, ``bias=False``, RMSNorm eps, tied-vs-untied LM head) that M4 +assembles into the full ``KimiLanguageModel``. Every builder matches the vLLM +DeepSeek-V3 reference: + - embedding / LM head: ``deepseek_v2.py`` ``DeepseekV2ForCausalLM`` (untied, + ``tie_word_embeddings=False``); + - dense MLP: ``DeepseekV2MLP`` = ``down_proj(SiluAndMul(gate_up_proj(x)))``, + ``bias=False``, silu-only; + - RMSNorm: standard Llama-style ``x * rsqrt(mean(x^2)+eps) * weight``. +""" +from __future__ import annotations + +from mstar.distributed.communication import CommGroup +from mstar.model.components import RMSNorm +from mstar.model.components.distributed import ( + ColumnParallelLinear, + ParallelGatedMLP, + VocabParallelEmbedding, +) +from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock +from mstar.model.kimi_k2_7.config import KimiK2Config + + +def build_embedding( + config: KimiK2Config, comm_group: CommGroup | None = None +) -> VocabParallelEmbedding: + """Token embedding ``[vocab, hidden]`` (row/vocab-parallel under TP).""" + return VocabParallelEmbedding( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + comm_group=comm_group, + padding_idx=config.pad_token_id, + ) + + +def build_lm_head( + config: KimiK2Config, comm_group: CommGroup | None = None +) -> ColumnParallelLinear: + """Untied LM head ``[hidden, vocab]`` (Kimi: ``tie_word_embeddings=False``). + + Column-parallel over vocab with ``gather_output=True`` so the sampler always + sees full ``[..., vocab]`` logits; a no-op all-gather at ``tp_size == 1``. + """ + return ColumnParallelLinear( + comm_group or CommGroup.trivial(), + input_size=config.hidden_size, + output_size=config.vocab_size, + bias=False, + gather_output=True, + ) + + +def build_rmsnorm(config: KimiK2Config) -> RMSNorm: + """Standard Llama-style RMSNorm (not Gemma's ``(1 + weight)`` variant).""" + return RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + +def build_dense_mlp( + config: KimiK2Config, comm_group: CommGroup | None = None +) -> ParallelGatedMLP: + """Dense SwiGLU MLP for the ``first_k_dense_replace`` early layers. + + Matches ``DeepseekV2MLP``: fused gate/up projection, ``silu(gate) * up``, + row-parallel down projection, ``bias=False``. Uses the full + ``intermediate_size`` (the MoE layers use ``moe_intermediate_size`` per + expert instead — that path is M2). + """ + return ParallelGatedMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + comm_group=comm_group, + activation=config.hidden_act, + bias=False, + ) + + +def is_moe_layer(config: KimiK2Config, layer_idx: int) -> bool: + """DeepSeek-V3 dense-vs-MoE layer selection. + + The first ``first_k_dense_replace`` layers are dense; thereafter every + ``moe_layer_freq``-th layer is MoE (``deepseek_v2.py`` decoder-layer ctor). + """ + return ( + layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + + +def build_moe_block( + config: KimiK2Config, comm_group: CommGroup | None = None +) -> KimiSparseMoeBlock: + """Fine-grained MoE block (routed experts + ungated shared expert).""" + return KimiSparseMoeBlock(config, comm_group=comm_group) + + +def build_mlp_for_layer( + config: KimiK2Config, layer_idx: int, comm_group: CommGroup | None = None +): + """Pick the layer's feed-forward: dense SwiGLU MLP or the MoE block. + + Returns a ``ParallelGatedMLP`` for the early dense layers, else a + ``KimiSparseMoeBlock``. Both expose the same ``(x) -> x`` interface, so the + decoder layer (M4) is agnostic to which it holds. + """ + if is_moe_layer(config, layer_idx): + return build_moe_block(config, comm_group=comm_group) + return build_dense_mlp(config, comm_group=comm_group) diff --git a/mstar/model/kimi_k2_7/components/moe.py b/mstar/model/kimi_k2_7/components/moe.py new file mode 100644 index 000000000..763909c11 --- /dev/null +++ b/mstar/model/kimi_k2_7/components/moe.py @@ -0,0 +1,258 @@ +"""Kimi-K2.7 / DeepSeek-V3 fine-grained MoE. + +mstar's ``model.components.moe`` router is softmax-only and its shared-expert +block gates the shared expert (Qwen-style). Kimi/DeepSeek-V3 needs a different +router and an *ungated* shared expert, so these live here (append, don't modify +the shared abstraction). The expert dispatch itself is reused verbatim — the +fused-expert GEMM (``fused_experts`` via ``model.components.moe._dispatch``) and +the ``(E, 2*moe_inter, hidden)`` / ``(E, hidden, moe_inter)`` fused param layout. + +Two pieces: + +* :class:`KimiMoEGate` — the router. sigmoid scoring + group-limited top-k + (``n_group`` / ``topk_group``) + ``noaux_tc`` per-expert + ``e_score_correction_bias`` (affects *selection* only; the combine weights come + from the raw sigmoid scores) + optional ``norm_topk_prob`` + a + ``routed_scaling_factor`` folded into the returned weights. Computed in fp32. + Exactly mirrors vLLM ``fused_moe/cpu_fused_moe.py::grouped_topk``. +* :class:`KimiSparseMoeBlock` — router + fused routed experts + ungated shared + expert. ``out = routed(scaled weights) + shared`` (the shared expert does *not* + get ``routed_scaling_factor``). Mirrors vLLM ``deepseek_v2.py::DeepseekV2MoE``. +""" +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn + +from mstar.distributed.communication import CommGroup +from mstar.model.components.distributed import ParallelGatedMLP +from mstar.model.components.moe import ( + _dispatch, + _down_proj_weight_loader, + _gate_up_weight_loader, +) +from mstar.model.kimi_k2_7.config import KimiK2Config + + +class KimiMoEGate(nn.Module): + """DeepSeek-V3 group-limited sigmoid router with ``noaux_tc`` bias. + + Args: + hidden_size: input hidden dim. + n_routed_experts: number of routed experts (``E``). + num_experts_per_tok: top-k experts per token. + n_group: number of expert groups (``E`` split into ``n_group`` contiguous + groups for group-limited routing). + topk_group: number of groups kept per token. + routed_scaling_factor: scale folded into the returned combine weights. + scoring_func: ``"sigmoid"`` (Kimi/DeepSeek-V3) or ``"softmax"``. + topk_method: ``"noaux_tc"`` enables the per-expert + ``e_score_correction_bias`` (selection-only). Anything else disables it. + norm_topk_prob: renormalize the top-k combine weights to sum to 1. + """ + + def __init__( + self, + hidden_size: int, + n_routed_experts: int, + num_experts_per_tok: int, + n_group: int, + topk_group: int, + routed_scaling_factor: float, + scoring_func: str = "sigmoid", + topk_method: str = "noaux_tc", + norm_topk_prob: bool = True, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.n_routed_experts = n_routed_experts + self.top_k = num_experts_per_tok + self.n_group = n_group + self.topk_group = topk_group + self.routed_scaling_factor = routed_scaling_factor + self.scoring_func = scoring_func + self.topk_method = topk_method + self.norm_topk_prob = norm_topk_prob + + # Router projection ``[E, hidden]`` (no bias), like DeepSeek ``MoEGate``. + self.weight = nn.Parameter(torch.zeros(n_routed_experts, hidden_size)) + if topk_method == "noaux_tc": + # Per-expert selection bias; fp32, added to scores for group/top-k + # selection but never to the combine weights. + self.e_score_correction_bias = nn.Parameter( + torch.zeros(n_routed_experts, dtype=torch.float32) + ) + else: + self.register_parameter("e_score_correction_bias", None) + + def forward( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Route tokens to experts. + + Returns: + topk_weights: ``(tokens, top_k)`` fp32 combine weights (renormalized + and scaled by ``routed_scaling_factor``). + topk_ids: ``(tokens, top_k)`` int64 expert indices. + """ + # Route in fp32 (DeepSeek runs the router in fp32 for stability). + h = hidden_states.reshape(-1, self.hidden_size).float() + gating = F.linear(h, self.weight.float()) # (T, E) + + if self.scoring_func == "sigmoid": + scores = gating.sigmoid() + elif self.scoring_func == "softmax": + scores = gating.softmax(dim=-1) + else: + raise ValueError(f"Unsupported scoring_func: {self.scoring_func!r}") + + num_token = scores.shape[0] + if self.e_score_correction_bias is not None: + # noaux_tc: bias-added scores drive group + expert *selection*; the + # raw sigmoid scores drive the combine weights. + original_scores = scores + scores = scores + self.e_score_correction_bias.unsqueeze(0) + group_scores = ( + scores.view(num_token, self.n_group, -1) + .topk(2, dim=-1)[0] + .sum(dim=-1) + ) # (T, n_group) + else: + original_scores = scores + group_scores = scores.view(num_token, self.n_group, -1).max(dim=-1).values + + group_idx = torch.topk( + group_scores, k=self.topk_group, dim=-1, sorted=False + )[1] # (T, topk_group) + group_mask = torch.zeros_like(group_scores) # (T, n_group) + group_mask.scatter_(1, group_idx, 1) + score_mask = ( + group_mask.unsqueeze(-1) + .expand(num_token, self.n_group, scores.shape[-1] // self.n_group) + .reshape(num_token, -1) + ) # (T, E) + masked_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) + + if self.e_score_correction_bias is not None: + topk_ids = torch.topk(masked_scores, k=self.top_k, dim=-1, sorted=False)[1] + topk_weights = original_scores.gather(1, topk_ids) + else: + topk_weights, topk_ids = torch.topk( + masked_scores, k=self.top_k, dim=-1, sorted=False + ) + + if self.norm_topk_prob: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + if self.routed_scaling_factor != 1.0: + topk_weights = topk_weights * self.routed_scaling_factor + + return topk_weights, topk_ids + + +class KimiSparseMoeBlock(nn.Module): + """DeepSeek-V3 MoE block: routed experts + ungated shared expert. + + ``out = routed(x) + shared(x)`` where ``routed`` dispatches the top-k experts + through the fused-expert GEMM with the router's (scaled) combine weights, and + ``shared`` is a plain dense SwiGLU MLP added ungated (no sigmoid gate, no + ``routed_scaling_factor``). + + Expert weights use the fused layout reused from ``model.components.moe``: + - ``experts.gate_up_proj``: ``(E, 2 * moe_intermediate_size, hidden)`` + - ``experts.down_proj``: ``(E, hidden, moe_intermediate_size)`` + """ + + def __init__( + self, config: KimiK2Config, comm_group: CommGroup | None = None + ) -> None: + super().__init__() + if comm_group is None: + comm_group = CommGroup.trivial() + self.comm_group = comm_group + self.hidden_size = config.hidden_size + self.num_experts = config.n_routed_experts + self.moe_intermediate_size = config.moe_intermediate_size + + self.gate = KimiMoEGate( + hidden_size=config.hidden_size, + n_routed_experts=config.n_routed_experts, + num_experts_per_tok=config.num_experts_per_tok, + n_group=config.n_group, + topk_group=config.topk_group, + routed_scaling_factor=config.routed_scaling_factor, + scoring_func=config.scoring_func, + topk_method=config.topk_method, + norm_topk_prob=config.norm_topk_prob, + ) + + self.experts = nn.Module() + self.experts.gate_up_proj = nn.Parameter( + torch.empty( + config.n_routed_experts, + 2 * config.moe_intermediate_size, + config.hidden_size, + ) + ) + self.experts.down_proj = nn.Parameter( + torch.empty( + config.n_routed_experts, + config.hidden_size, + config.moe_intermediate_size, + ) + ) + # The fused expert params are plain nn.Parameters, so they carry no + # per-shard ``weight_loader`` by default. The M5 stacked-param rules route + # each checkpoint expert via a ``"gate:N"/"up:N"/"down:N"`` shard id, so we + # attach the same fused-expert loaders ``ParallelSparseMoeBlock`` uses. + # Experts are held full-size here (no expert/TP sharding yet — TODO(M6)), + # hence ``tp_rank=0, tp_size=1`` and ``full_inter == moe_intermediate_size``. + self._attach_expert_weight_loaders() + + # Ungated shared expert: a dense SwiGLU MLP with the shared intermediate + # size (``moe_intermediate_size * n_shared_experts``). + self.shared_expert = ParallelGatedMLP( + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size * config.n_shared_experts, + comm_group=comm_group, + activation=config.hidden_act, + bias=False, + ) + + def _attach_expert_weight_loaders(self) -> None: + """Give the fused expert params their per-shard ``weight_loader``. + + Mirrors ``ParallelSparseMoeBlock._attach_weight_loaders``. Re-run after + every ``_apply`` (``.to(dtype)`` / ``to_empty(device)`` rebuild the + Parameter objects and drop the attribute), so weights load correctly + through the meta -> to_empty -> load path. + """ + from functools import partial + + self.experts.gate_up_proj.weight_loader = partial( + _gate_up_weight_loader, 0, 1, self.moe_intermediate_size, + ) + self.experts.down_proj.weight_loader = partial( + _down_proj_weight_loader, 0, 1, self.moe_intermediate_size, + ) + + def _apply(self, fn, recurse=True): + result = super()._apply(fn, recurse=recurse) + self._attach_expert_weight_loaders() + return result + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_shape = hidden_states.shape + flat = hidden_states.view(-1, self.hidden_size).contiguous() + + topk_weights, topk_ids = self.gate(flat) + routed = _dispatch( + flat, + self.experts.gate_up_proj, + self.experts.down_proj, + self.num_experts, + topk_ids, + topk_weights.to(flat.dtype), + ) + shared = self.shared_expert(flat) + return (routed + shared).view(input_shape) diff --git a/mstar/model/kimi_k2_7/components/rope.py b/mstar/model/kimi_k2_7/components/rope.py new file mode 100644 index 000000000..3cde1b4f9 --- /dev/null +++ b/mstar/model/kimi_k2_7/components/rope.py @@ -0,0 +1,146 @@ +"""deepseek_yarn RoPE for Kimi-K2.7 / DeepSeek-V3 MLA. + +MLA rotates only the decoupled ``qk_rope_head_dim`` slice of q/k, with YARN +(NTK-by-parts) frequency scaling and an ``mscale`` amplitude on cos/sin. mstar's +``cache_manager.apply_rope`` (FlashInfer) does not implement YARN, so this is a +standalone rotary module the MLA attention applies itself (analogous to how +Qwen3-Omni applies its 3D MRoPE outside the cache handle). + +Style is **interleaved / GPT-J** (``is_neox_style=False`` in DeepSeek): cos/sin +are ``repeat_interleave(2)`` and adjacent even/odd pairs are rotated. Mirrors +vLLM ``layers/rotary_embedding/deepseek_scaling_rope.py::DeepseekScalingRotaryEmbedding`` +and the YARN helpers in ``rotary_embedding/common.py``. + +Two ``mscale`` values (both use the 2-arg ``yarn_get_mscale``): + - **amplitude** on cos/sin (here): ``get_mscale(f, mscale) / get_mscale(f, mscale_all_dim) * attn_factor``. + - **softmax-scale boost** (in the attention, not here): ``get_mscale(f, mscale_all_dim) ** 2``. +""" +from __future__ import annotations + +import math + +import torch +from torch import nn + + +def yarn_get_mscale(scale: float = 1.0, mscale: float = 1.0) -> float: + """DeepSeek 2-arg mscale (deepseek_v2.py:428 / deepseek_scaling_rope.py:20).""" + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +def rotate_gptj(x: torch.Tensor) -> torch.Tensor: + """Interleaved (GPT-J) rotate: pairs ``x[..., ::2]`` / ``x[..., 1::2]`` + (vLLM ``common.py::rotate_gptj``).""" + x1 = x[..., ::2] + x2 = x[..., 1::2] + return torch.stack((-x2, x1), dim=-1).flatten(-2) + + +def _yarn_find_correction_dim(num_rotations, dim, base, max_pos) -> float: + return (dim * math.log(max_pos / (num_rotations * 2 * math.pi))) / (2 * math.log(base)) + + +def _yarn_find_correction_range(low_rot, high_rot, dim, base, max_pos) -> tuple[int, int]: + low = math.floor(_yarn_find_correction_dim(low_rot, dim, base, max_pos)) + high = math.ceil(_yarn_find_correction_dim(high_rot, dim, base, max_pos)) + return max(low, 0), min(high, dim - 1) + + +def _yarn_linear_ramp_mask(low, high, dim, dtype) -> torch.Tensor: + if low == high: + high += 0.001 # avoid singularity + ramp = (torch.arange(dim, dtype=dtype) - low) / (high - low) + return torch.clamp(ramp, 0, 1) + + +class KimiYarnRotaryEmbedding(nn.Module): + """deepseek_yarn rotary embedding over the ``rotary_dim`` (=qk_rope_head_dim) slice.""" + + def __init__( + self, + rotary_dim: int, + base: float, + factor: float, + original_max_position_embeddings: int, + beta_fast: float = 32, + beta_slow: float = 1, + mscale: float = 1.0, + mscale_all_dim: float = 0.0, + extrapolation_factor: float = 1.0, + attn_factor: float = 1.0, + ) -> None: + super().__init__() + self.rotary_dim = rotary_dim + + # ``inv_freq`` is NOT a registered buffer. Buffers computed in ``__init__`` + # do not survive the production ``meta`` build -> ``to_empty(device)`` -> + # ``load_weights`` path: ``to_empty`` allocates uninitialized memory and + # never re-runs ``__init__``, and ``inv_freq`` is not in the checkpoint + # (it's derived, skipped by the loader) — so a buffer would be left as + # garbage after loading, silently corrupting YARN RoPE. Instead keep the + # scalar recipe and compute ``inv_freq`` lazily in fp32 on the target + # device (also keeps it fp32 under a bf16 model, matching DeepSeek, rather + # than being downcast by ``model.to(bf16)``). + self._inv_freq_args = ( + rotary_dim, base, factor, original_max_position_embeddings, + beta_fast, beta_slow, extrapolation_factor, + ) + self._inv_freq_cache: torch.Tensor | None = None + + # cos/sin amplitude (deepseek_scaling_rope.py:56-60). + self.mscale = float( + yarn_get_mscale(factor, mscale) + / yarn_get_mscale(factor, mscale_all_dim) + * attn_factor + ) + + def _get_inv_freq(self, device: torch.device) -> torch.Tensor: + """Return the fp32 ``inv_freq`` for ``device``, computing + caching once.""" + cached = self._inv_freq_cache + if cached is None or cached.device != device: + cached = self._compute_inv_freq(*self._inv_freq_args).to(device=device) + self._inv_freq_cache = cached + return cached + + @staticmethod + def _compute_inv_freq( + rotary_dim, base, factor, max_pos, beta_fast, beta_slow, extrapolation_factor, + ) -> torch.Tensor: + pos_freqs = base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim) + inv_freq_extrapolation = 1.0 / pos_freqs + inv_freq_interpolation = 1.0 / (factor * pos_freqs) + + low, high = _yarn_find_correction_range( + beta_fast, beta_slow, rotary_dim, base, max_pos + ) + inv_freq_mask = ( + 1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float) + ) * extrapolation_factor + return ( + inv_freq_interpolation * (1 - inv_freq_mask) + + inv_freq_extrapolation * inv_freq_mask + ) + + def forward( + self, position_ids: torch.Tensor, q_pe: torch.Tensor, k_pe: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Rotate the pe slices. + + Args: + position_ids: ``(tokens,)`` int positions. + q_pe: ``(tokens, num_heads, rotary_dim)``. + k_pe: ``(tokens, 1, rotary_dim)`` (shared MQA rope key). + Returns: + rotated ``(q_pe, k_pe)`` in the input dtypes. + """ + inv_freq = self._get_inv_freq(position_ids.device) + freqs = torch.outer(position_ids.float(), inv_freq) # (T, rotary_dim/2) + cos = (freqs.cos() * self.mscale).repeat_interleave(2, dim=-1).unsqueeze(-2) + sin = (freqs.sin() * self.mscale).repeat_interleave(2, dim=-1).unsqueeze(-2) + + q32, k32 = q_pe.float(), k_pe.float() + q_rot = q32 * cos + rotate_gptj(q32) * sin + k_rot = k32 * cos + rotate_gptj(k32) * sin + return q_rot.to(q_pe.dtype), k_rot.to(k_pe.dtype) diff --git a/mstar/model/kimi_k2_7/config.py b/mstar/model/kimi_k2_7/config.py new file mode 100644 index 000000000..12267cffe --- /dev/null +++ b/mstar/model/kimi_k2_7/config.py @@ -0,0 +1,151 @@ +"""Configuration dataclass for Kimi-K2.7 (text backbone). + +Kimi-K2.7's text architecture *is* DeepSeek-V3 — vLLM serves it as +``DeepseekV3ForCausalLM`` (``model_type: "kimi_k2"`` maps to +``DeepseekV3Config``). This dataclass therefore carries the full DeepSeek-V3 +field set: MLA latent dims, fine-grained sigmoid-routed MoE grouping, and +``deepseek_yarn`` RoPE. Only a handful of these fields are read by the M0 +scaffold (``num_hidden_layers``, the head dims, ``vocab_size``, +``max_position_embeddings``); the rest are declared now so this stays the single +source of truth for the later milestones (MoE router, MLA attention, weights). + +The full-size defaults below are the real values from the +``moonshotai/Kimi-K2.7-Code`` HF ``config.json`` (``model_type: "kimi_k2"`` → +``DeepseekV3Config``), confirmed field by field against the published checkpoint +config (the weights themselves are not present in this workspace). M0 does not +depend on the full-size values being exact — the modular tests build from +:meth:`KimiK2Config.reduced`, a tiny self-consistent config. +""" +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass +class KimiK2Config: + # -- Core transformer dims -------------------------------------------- + vocab_size: int = 163840 + hidden_size: int = 7168 + intermediate_size: int = 18432 # dense-FFN size (first_k_dense_replace layers) + num_hidden_layers: int = 61 + num_attention_heads: int = 64 + num_key_value_heads: int = 64 # MLA has no separate KV heads; kept for HF parity + rms_norm_eps: float = 1e-5 # from config.json (Kimi uses 1e-5, not DeepSeek-V3's 1e-6) + max_position_embeddings: int = 262144 + tie_word_embeddings: bool = False + hidden_act: str = "silu" + + # -- MLA (Multi-head Latent Attention) latent dims -------------------- + # Query is compressed to ``q_lora_rank`` then projected up to + # ``num_attention_heads * qk_head_dim``; K/V share a ``kv_lora_rank`` latent + # plus a decoupled ``qk_rope_head_dim`` RoPE slice. Per-head query/key dim is + # ``qk_nope_head_dim + qk_rope_head_dim``; value head dim differs. + q_lora_rank: int = 1536 + kv_lora_rank: int = 512 + qk_nope_head_dim: int = 128 + qk_rope_head_dim: int = 64 + v_head_dim: int = 128 + + # -- Fine-grained MoE (sigmoid router, group-limited top-k, noaux_tc) -- + n_routed_experts: int = 384 # from config.json + n_shared_experts: int = 1 + num_experts_per_tok: int = 8 # top-k + moe_intermediate_size: int = 2048 + n_group: int = 1 # from config.json + topk_group: int = 1 # from config.json (groups kept by group-limited routing) + routed_scaling_factor: float = 2.827 # from config.json + scoring_func: str = "sigmoid" # DeepSeek-V3/Kimi: sigmoid (not softmax) + topk_method: str = "noaux_tc" # per-expert e_score_correction_bias + norm_topk_prob: bool = True + first_k_dense_replace: int = 1 # first N layers are dense, rest are MoE + moe_layer_freq: int = 1 + + # -- deepseek_yarn RoPE ------------------------------------------------ + rope_theta: float = 50000.0 # from config.json + rope_scaling: dict = field(default_factory=lambda: { + # from config.json (HF key is "type": "yarn"; mstar's internal id for the + # DeepSeek/Kimi variant is "deepseek_yarn"). factor=64 yields the 262144 + # context (4096 * 64). K2.7-Code keeps beta_fast=32 (some other Kimi + # checkpoints set beta_fast=1). mscale == mscale_all_dim == 1.0. + "rope_type": "deepseek_yarn", + "factor": 64.0, + "original_max_position_embeddings": 4096, + "beta_fast": 32.0, + "beta_slow": 1.0, + "mscale": 1.0, + "mscale_all_dim": 1.0, + }) + + # -- Special tokens / generation defaults ----------------------------- + bos_token_id: int = 163584 # from config.json + eos_token_id: int = 163586 # from config.json + pad_token_id: int = 163839 # from config.json + temperature: float = 1.0 + top_p: float = 1.0 + ignore_eos: bool = False + + # -- MTP (multi-token prediction) — deferred, declared for completeness - + num_nextn_predict_layers: int = 0 + + # --------------------------------------------------------------------- + # Derived dims (read by get_kv_cache_config / attention) + # --------------------------------------------------------------------- + @property + def qk_head_dim(self) -> int: + """Per-head query/key dim: nope + decoupled-rope slice (e.g. 128+64=192).""" + return self.qk_nope_head_dim + self.qk_rope_head_dim + + @property + def padded_head_dim(self) -> int: + """Head dim the naive-MLA q/k/v are zero-padded to for the paged cache. + + FlashInfer's SM90 (Hopper) prefill kernel ``static_assert``s + ``head_dim_vo ∈ {64, 128, 256}`` (M4 finding), so it will not JIT-build for + the real ``qk_head_dim=192`` or the reduced ``qk_head_dim=24``. The + correctness-first mitigation (M6) pads q/k (from ``qk_head_dim``) and v + (from ``v_head_dim``) up to the smallest supported dim ``>= qk_head_dim``, + runs the paged attention there, and slices the output back to + ``v_head_dim`` — compensating the softmax scale (see + ``KimiMLAAttention.softmax_scale_boost``). Real Kimi 192 -> 256; reduced + 24 -> 64. Weight-absorbed MLA (which avoids the pad) is deferred to perf. + """ + for supported in (64, 128, 256): + if supported >= self.qk_head_dim: + return supported + raise ValueError( + f"qk_head_dim={self.qk_head_dim} exceeds the largest FlashInfer SM90 " + "head_dim (256); the naive-MLA pad mitigation cannot cover it." + ) + + @property + def num_dense_layers(self) -> int: + return min(self.first_k_dense_replace, self.num_hidden_layers) + + @classmethod + def reduced(cls) -> "KimiK2Config": + """A tiny, self-consistent config for CPU/dummy-mode modular tests and + reduced-config golden runs. Keeps the *shape* of Kimi (MLA split heads, + grouped MoE, one dense layer) while being small enough to run without + the 1T checkpoint. + """ + return cls( + vocab_size=256, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + max_position_embeddings=512, + q_lora_rank=48, + kv_lora_rank=32, + qk_nope_head_dim=16, + qk_rope_head_dim=8, + v_head_dim=16, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + moe_intermediate_size=64, + n_group=1, + topk_group=1, + first_k_dense_replace=1, + ) diff --git a/mstar/model/kimi_k2_7/kimi_model.py b/mstar/model/kimi_k2_7/kimi_model.py new file mode 100644 index 000000000..3149a455a --- /dev/null +++ b/mstar/model/kimi_k2_7/kimi_model.py @@ -0,0 +1,386 @@ +"""KimiK2Model: M* Model contract for Kimi-K2.7 (text backbone). + +Kimi-K2.7's text path is DeepSeek-V3 (``model_type: "kimi_k2"`` → +``DeepseekV3ForCausalLM``). This is the **M0 scaffold**: it declares the full +serving plumbing — the graph (``prefill`` + ``decode`` Loop), the single +``KV_CACHE`` LLM node, the KV-cache dims, and the prefill→decode→done state +machine — with **no GPU compute**. ``get_submodule`` returns ``None`` (dummy +mode), so ``pytest test/modular/`` exercises the graph/walk/engine-routing +machinery in isolation, exactly as ``docs/adding_models.rst`` prescribes for a +new model before touching weights. + +Structurally this mirrors Orpheus's LLM partition (the smallest complete LLM in +the tree) minus the async SNAC partition: Kimi text-only is a single ``default`` +partition, so it inherits ``Model.get_partitions`` / ``get_partition_topology`` +and only implements the abstract surface. + +Later milestones fill in the real compute (M2 MoE router, M3 MLA attention, M5 +weights, M6 the ``KimiLLMSubmodule`` build in ``get_submodule``); none of them +change the contract declared here. +""" +from __future__ import annotations + +import logging + +import torch + +from mstar.communication.tensors import NameToTensorList +from mstar.conductor.request_info import ( + CurrentForwardConductorMetadata, + StreamingConnectionState, +) +from mstar.engine.base import EngineType +from mstar.engine.kv_cache_engine import KVCacheConfig +from mstar.graph.base import GraphEdge, GraphNode, GraphSection, Loop, TensorPointerInfo +from mstar.graph.special_destinations import EMIT_TO_CLIENT +from mstar.model.base import ForwardPassArgs, Model +from mstar.model.kimi_k2_7.config import KimiK2Config +from mstar.model.submodule_base import NodeSubmodule +from mstar.utils.sampling import SamplingConfig + +logger = logging.getLogger(__name__) + +LLM_NODE = "LLM" +DECODE_LOOP = "decode_loop" + + +def _resolve_local_hf_snapshot(repo_id: str, cache_dir: str | None = None) -> str: + """Resolve an HF repo id to a local snapshot dir (mirrors OrpheusModel).""" + from pathlib import Path + + from huggingface_hub import snapshot_download + + try: + local_dir = snapshot_download( + repo_id=repo_id, cache_dir=cache_dir, local_files_only=False, + ) + except Exception as e: # noqa: BLE001 + logger.warning("Error downloading %r from huggingface: %s", repo_id, e) + return repo_id + return str(Path(local_dir)) + + +class KimiK2Model(Model): + """Kimi-K2.7 text backbone (DeepSeek-V3 architecture).""" + + def __init__( + self, + model_path_hf: str, + cache_dir: str | None = None, + **kwargs, + ): + self.cache_dir = cache_dir + self.model_path_hf = model_path_hf + self.config = KimiK2Config() + # Tokenizer is loaded lazily: the modular (dummy-mode) tests build the + # model via ``object.__new__`` and never call ``__init__``, so we avoid + # forcing a network/tokenizer dependency into the scaffold path. + self._tokenizer = None + self._submodule_cache: dict[str, NodeSubmodule | None] = {} + + @property + def tokenizer(self): + if self._tokenizer is None: + from transformers import AutoTokenizer + + self._tokenizer = AutoTokenizer.from_pretrained( + self.model_path_hf, + cache_dir=self.cache_dir, + trust_remote_code=True, + ) + return self._tokenizer + + # ------------------------------------------------------------------- + # Model ABC: KV cache config + # ------------------------------------------------------------------- + + def get_kv_cache_config(self) -> list[KVCacheConfig]: + # Naive/materialized MLA (the first-pass port, per CLAUDE.md): the latent + # is projected up to full per-head K/V and broadcast to every query head, + # so from the paged cache's ``[tokens, heads, head_dim]`` point of view + # there are ``num_attention_heads`` KV heads. K/V are stored at + # ``padded_head_dim`` — the naive path zero-pads q/k (from ``qk_head_dim``, + # e.g. 192) and v (from ``v_head_dim``) up to the smallest FlashInfer-SM90 + # supported head_dim >= qk_head_dim (256 real, 64 reduced), because the + # Hopper prefill kernel static_asserts head_dim_vo in {64,128,256} (M4 + # finding). The attention output is sliced back to ``v_head_dim`` in the + # submodule (M3/M6). This trades cache size for not needing a weight-absorb + # path in the engine (deferred to perf). + return [KVCacheConfig( + num_layers=self.config.num_hidden_layers, + num_kv_heads=self.config.num_attention_heads, + head_dim=self.config.padded_head_dim, + max_seq_len=self.config.max_position_embeddings, + num_qo_heads=self.config.num_attention_heads, + )] + + # ------------------------------------------------------------------- + # Model ABC: node engine types + # ------------------------------------------------------------------- + + def get_node_engine_types(self) -> dict[str, EngineType]: + return {LLM_NODE: EngineType.KV_CACHE} + + # ------------------------------------------------------------------- + # Model ABC: graph walk definitions + # ------------------------------------------------------------------- + + def get_graph_walk_graphs(self) -> dict[str, GraphSection]: + # prefill: embed the prompt, fill the KV cache, sample + emit the first + # token. ``persist=True`` keeps that token at the conductor so the decode + # walk can pick it up as its first ``text_inputs``. + prefill = GraphNode( + name=LLM_NODE, + input_names=["text_inputs"], + outputs=[ + GraphEdge( + next_node=EMIT_TO_CLIENT, + name="new_token", + output_modality="text", + conductor_new_token=True, + persist=True, + ), + ], + ) + + # decode: autoregressive Loop. Each step emits the new token to the + # client and feeds it back as the next step's ``text_inputs``. The Loop + # stops via the submodule's ``check_stop`` (EOS / max tokens); ``max_iters`` + # is the hard cap. + decode = Loop( + name=DECODE_LOOP, + section=GraphNode( + name=LLM_NODE, + input_names=["text_inputs"], + outputs=[ + GraphEdge( + next_node=EMIT_TO_CLIENT, + name="new_token", + output_modality="text", + conductor_new_token=True, + ), + GraphEdge( + next_node=LLM_NODE, + name="text_inputs", + ), + ], + ), + max_iters=self.get_max_output_tokens(), + outputs=[], + ) + + return dict(prefill=prefill, decode=decode) + + # ------------------------------------------------------------------- + # Model ABC: forward pass args (single "default" partition) + # ------------------------------------------------------------------- + + def get_initial_forward_pass_args( + self, + partition_name: str, + input_modalities: list[str], + output_modalities: list[str], + input_signals: dict[str, list[TensorPointerInfo]], + model_kwargs: dict | None = None, + ) -> ForwardPassArgs: + full_metadata = CurrentForwardConductorMetadata( + input_modalities=input_modalities, + output_modalities=output_modalities, + graph_walk="prefill", + is_prefill=True, + ) + + graph_edge = GraphEdge(next_node=LLM_NODE, name="text_inputs") + graph_edge.tensor_info = input_signals.get("text_inputs", []) + inputs = [graph_edge] + unpersist_tensors = sum([inp.tensor_info for inp in inputs], start=[]) + + return ForwardPassArgs( + full_metadata=full_metadata, + inputs=inputs, + unpersist_tensors=unpersist_tensors, + step_metadata={"is_prefill": True}, + ) + + def get_partition_forward_pass_args( + self, + partition_name: str, + partition_metadata: CurrentForwardConductorMetadata, + persist_signals: dict[str, list[TensorPointerInfo]], + incoming_connections: list[StreamingConnectionState] | None = None, + ) -> ForwardPassArgs: + """Drive the prefill → decode → done state machine. + + Called by the conductor after each completed walk. Prefill transitions to + the decode walk (feeding the persisted first token as ``text_inputs``); + once the decode walk completes (its Loop stopped via ``check_stop``), the + request is done. The per-token decode iteration is driven inside the + graph Loop, not by repeated calls here. + """ + metadata = partition_metadata + request_done = False + + if metadata.is_prefill: + metadata.is_prefill = False + metadata.graph_walk = "decode" + elif metadata.graph_walk == "decode": + request_done = True + metadata.kwargs["decode_finished"] = True + + if request_done: + return ForwardPassArgs( + full_metadata=metadata, + inputs=[], + unpersist_tensors=[], + request_done=True, + ) + + graph_edge = GraphEdge(next_node=LLM_NODE, name="text_inputs") + graph_edge.tensor_info = persist_signals.get("new_token", []) + inputs = [graph_edge] + unpersist_tensors = sum([inp.tensor_info for inp in inputs], start=[]) + + return ForwardPassArgs( + full_metadata=metadata, + inputs=inputs, + unpersist_tensors=unpersist_tensors, + step_metadata={"is_prefill": metadata.is_prefill}, + ) + + # ------------------------------------------------------------------- + # Model ABC: prompt processing + # ------------------------------------------------------------------- + + def process_prompt( + self, + prompt: str | None, + input_modalities: list[str], + output_modalities: list[str], + tensors: NameToTensorList | None = None, + **kwargs, + ) -> NameToTensorList: + # Text-only for M0; raw multimodal tensors (MoonViT) are a later milestone. + if prompt is None: + return {} + input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids[0] + return {"text_inputs": [input_ids]} + + def get_sampling_config( + self, + node_name: str, + model_kwargs: dict | None = None, + ) -> SamplingConfig | None: + model_kwargs = model_kwargs or {} + return SamplingConfig( + vocab_size=self.config.vocab_size, + temperature=model_kwargs.get("temperature", self.config.temperature), + top_p=model_kwargs.get("top_p", self.config.top_p), + ignore_eos=model_kwargs.get("ignore_eos", self.config.ignore_eos), + ) + + # ------------------------------------------------------------------- + # Model ABC: postprocess + # ------------------------------------------------------------------- + + def postprocess( + self, + output: torch.Tensor, + modality: str, + request_kwargs: dict | None = None, + ) -> bytes: + if modality == "text": + token_ids = output.tolist() if output.numel() else [] + text = self.tokenizer.decode(token_ids, skip_special_tokens=True) + return text.encode("utf-8") + raise ValueError(f"Unsupported modality for Kimi-K2.7: {modality!r}") + + # ------------------------------------------------------------------- + # Model ABC: sharding + # ------------------------------------------------------------------- + + def get_default_sharding_config(self): + from mstar.distributed.base import ShardingConfig + + # Kimi is a 1T model — real serving is TP8 / multi-node. The LLM node is + # the tensor-parallel node; the per-node degree comes from the config + # YAML's ``node_groups`` (M6), not from the model code. + return ShardingConfig(groups=[], tp_enabled_nodes={LLM_NODE}, shard_dim={}) + + # ------------------------------------------------------------------- + # Model ABC: submodule loading + # ------------------------------------------------------------------- + + def get_submodule( + self, + node_name: str, + device: str = "cpu", + tp_group=None, + autocast_dtype: torch.dtype | None = None, + sp_group=None, + ) -> NodeSubmodule | None: + if node_name in self._submodule_cache: + return self._submodule_cache[node_name] + submodule = self._create_submodule( + node_name, device, tp_group=tp_group, autocast_dtype=autocast_dtype, + ) + self._submodule_cache[node_name] = submodule + return submodule + + def _create_submodule( + self, + node_name: str, + device: str, + tp_group=None, + autocast_dtype: torch.dtype | None = None, + ) -> NodeSubmodule | None: + if node_name != LLM_NODE: + return None + + source = self._resolve_checkpoint() + if source is None: + # Dummy mode: no checkpoint resolvable (e.g. the modular graph tests + # build the model via object.__new__ with no model_path_hf). Returning + # None lets pytest test/modular/ validate the graph/walks/engine-routing + # without a GPU or weights, per docs/adding_models.rst. + logger.info( + "KimiK2Model: no checkpoint resolved for node %r — dummy mode (None).", + node_name, + ) + return None + + # Real build, mirroring OrpheusModel._create_llm_submodule: construct on the + # meta device (no allocation), cast to the target dtype on meta (so to_empty + # allocates directly in bf16, not fp32-then-downcast), materialise storage, + # then run the M5 HF loader (remap + fused-expert stacked rules). This is the + # path the M5 rope-buffer bug would have bitten — inv_freq is lazy so it does + # not survive as garbage. + from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM + from mstar.model.kimi_k2_7.submodules import KimiLLMSubmodule + from mstar.model.loader import load_weights + + with torch.device("meta"): + language_model = KimiForCausalLM(self.config, comm_group=tp_group) + if autocast_dtype is not None: + language_model = language_model.to(autocast_dtype) + language_model.to_empty(device=device) + load_weights(language_model, source, device=device) + language_model.eval() + + logger.info("Successfully loaded Kimi-K2.7 submodule for %s", node_name) + return KimiLLMSubmodule(language_model=language_model, config=self.config) + + def _resolve_checkpoint(self) -> str | None: + """Resolve the HF checkpoint source, or None for dummy mode. + + A local directory / file (e.g. a reduced synthetic checkpoint) is used + as-is; otherwise the HF repo id is snapshot-downloaded. Returns None when + no ``model_path_hf`` is set (dummy-mode graph tests). + """ + from pathlib import Path + + path = getattr(self, "model_path_hf", None) + if not path: + return None + if Path(path).exists(): + return str(path) + return _resolve_local_hf_snapshot(path, cache_dir=getattr(self, "cache_dir", None)) diff --git a/mstar/model/kimi_k2_7/submodules.py b/mstar/model/kimi_k2_7/submodules.py new file mode 100644 index 000000000..015ce5e5a --- /dev/null +++ b/mstar/model/kimi_k2_7/submodules.py @@ -0,0 +1,269 @@ +"""Submodules for Kimi-K2.7 (text backbone). + +M6: the real :class:`KimiLLMSubmodule` — the ``ARNodeSubmodule`` that drives the +DeepSeek-V3 text backbone (MLA attention over the paged cache + fine-grained +sigmoid-routed MoE) through the engine's ``prepare_inputs -> preprocess -> +forward/forward_batched -> postprocess -> check_stop`` lifecycle for the +``prefill`` and ``decode`` Loop walks. + +Structurally this mirrors ``OrpheusLLMSubmodule`` (the smallest complete LLM in +the tree) with one addition: the naive MLA applies YARN RoPE itself over the +decoupled ``qk_rope`` slice, so ``preprocess`` builds per-token ``position_ids`` +(the same positions ``plan_rope`` uses) and threads them into the forward — +analogous to how Qwen3-Omni threads its 3D-MRoPE cos/sin through preprocess. The +sampling/EOS/logits contract is identical to Orpheus: the non-batched ``forward`` +returns last-token ``logits`` (the KV-cache engine samples them into +``new_token``); ``forward_batched`` samples inside the forward and returns +``new_token`` per request. +""" +from __future__ import annotations + +from typing import Any + +import torch +from torch import nn + +from mstar.communication.tensors import NameToTensorList +from mstar.conductor.request_info import CurrentForwardPassInfo +from mstar.engine.base import NodeBatch +from mstar.engine.cache_manager import BatchedCacheManager +from mstar.engine.cuda_graph_config import FlashInferPackedCudaGraphConfig +from mstar.engine.cuda_graph_runner import BasicBatchedCudaGraphConfig +from mstar.engine.kv_store import PositionInfo +from mstar.model.kimi_k2_7.config import KimiK2Config +from mstar.model.submodule_base import ( + ARNodeInputs, + ARNodeSubmodule, + ModelInputsFromEngine, + NodeInputs, +) +from mstar.utils.sampling import Sampler + +_MAIN = "main" + + +class KimiLLMSubmodule(ARNodeSubmodule): + """Autoregressive Kimi/DeepSeek-V3 text backbone (prefill + decode). + + Dispatches on ``graph_walk``: + - ``prefill``: embed the prompt, fill the KV cache, sample the first token; + - ``decode``: embed the previous token, generate the next token. + """ + + def __init__(self, language_model: nn.Module, config: KimiK2Config): + super().__init__() + self.language_model = language_model # KimiForCausalLM + self.lm_head = language_model.lm_head + self.config = config + + # -- CUDA-graph capture buckets (mirror OrpheusLLMSubmodule) ------------ + PREFILL_TOKEN_BUCKETS = [32, 64, 128, 256, 512, 1024] + PREFILL_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16] + + def _build_prefill_packed( + self, num_tokens: int, device: torch.device, + ) -> dict[str, torch.Tensor]: + """Tensor-only post-``preprocess`` packed dict for prefill capture. + + Mirrors ``preprocess``'s tensor outputs: the packed ``(num_tokens,)`` long + ``input_ids`` and the ``(num_tokens,)`` long ``position_ids`` the MLA YARN + RoPE reads. Both are interned as static buffers; at replay the runner + copies the real ``preprocess`` output into them. + """ + return { + "input_ids": torch.zeros((num_tokens,), dtype=torch.long, device=device), + "position_ids": torch.arange(num_tokens, dtype=torch.long, device=device), + } + + def get_cuda_graph_configs( + self, device: torch.device, tp_world_size: int = 1, + ) -> list[BasicBatchedCudaGraphConfig | FlashInferPackedCudaGraphConfig]: + """Decode (per-bs) + prefill (per-token-bucket) captures. + + The YARN-rope path the MLA runs is pure tensor compute (``outer`` + cos/sin + + interleaved rotate) reading ``position_ids`` from a static buffer, so it + captures like Qwen3-Omni's cos/sin-threaded prefill: decode re-runs + ``preprocess`` at replay and copies the packed ``input_ids`` / + ``position_ids`` into the interned buffers; prefill uses the packed dict + above. ``inv_freq`` is lazily cached on first (warmup) forward, so its + address is stable across replay. + """ + prefill_packed = { + num_tokens: self._build_prefill_packed(num_tokens, device) + for num_tokens in self.PREFILL_TOKEN_BUCKETS + } + return [ + BasicBatchedCudaGraphConfig( + capture_graph_walk="decode", + requires_cfg=False, + labels=[_MAIN], + single_request_inputs=ARNodeInputs( + input_ids=torch.zeros(1, dtype=torch.long, device=device), + input_seq_len=1, + ), + ), + FlashInferPackedCudaGraphConfig( + capture_graph_walk="prefill", + replay_graph_walks=["prefill"], + packed_seq_len_to_inputs=prefill_packed, + requires_cfg=False, + labels=[_MAIN], + compile=True, + causal_attention=True, + capture_batch_sizes=self.PREFILL_CAPTURE_BATCH_SIZES, + ), + ] + + # -- lifecycle --------------------------------------------------------- + + def prepare_inputs( + self, + graph_walk: str, + fwd_info: CurrentForwardPassInfo, + inputs: NameToTensorList, + pos_info: dict[str, PositionInfo] = {}, + **kwargs, + ) -> ARNodeInputs: + # Cheap host-side: the prompt ids (prefill) or the previous token (decode) + # arrive under the "text_inputs" edge (see KimiK2Model.get_graph_walk_graphs). + text_inputs = inputs["text_inputs"][0] + return ARNodeInputs( + input_ids=text_inputs, + input_seq_len=text_inputs.shape[0], + ) + + def preprocess( + self, + graph_walk: str, + engine_inputs: ModelInputsFromEngine, + inputs: list[ARNodeInputs], + ) -> dict[str, torch.Tensor | Any]: + cache_manager = engine_inputs.cache_manager + seq_lens = [inp.input_seq_len for inp in inputs] + + # Plan attention + rope for the main cache label (CUDA-graph incompatible, + # so it happens here in preprocess, not in forward). + cache_manager.set_active_label(_MAIN) + cache_manager.plan_attention(seq_lens=seq_lens, is_causal=True, label=_MAIN) + cache_manager.plan_rope(seq_lens=seq_lens, pos_ids=None, label=_MAIN) + + # Build the per-token YARN position_ids for the MLA rope. These are exactly + # the positions plan_rope uses: each request's position_id_start (advanced + # once per forward by KimiLanguageModel's advance_seq_lens) plus its span. + # request_ids order matches the inputs order (both are batch order), so the + # concatenated input_ids and position_ids stay token-aligned. + device = self.get_device() + pos_ids_list: list[int] = [] + for rid, sl in zip(cache_manager.request_ids, seq_lens, strict=True): + start = cache_manager._get_state(rid, _MAIN).position_id_start + pos_ids_list.extend(range(start, start + sl)) + position_ids = torch.tensor(pos_ids_list, dtype=torch.long, device=device) + + return { + "input_ids": torch.cat([inp.input_ids for inp in inputs]), + "position_ids": position_ids, + } + + def _hidden( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + cache_handle: BatchedCacheManager, + ) -> torch.Tensor: + """Token ids -> final hidden states (embed -> layers -> norm). + + ``KimiLanguageModel.forward`` embeds ``input_ids``, runs the decoder stack + (per-layer ``set_layer_idx``, MLA YARN rope from ``position_ids``), calls + ``advance_seq_lens`` once, and returns the normed hidden states. + """ + return self.language_model.model(input_ids, cache_handle, position_ids) + + def forward( + self, + graph_walk: str, + engine_inputs: ModelInputsFromEngine, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + **kwargs, + ) -> NameToTensorList: + """Non-batched forward: return the last token's logits for the engine to + sample (prefill: last prompt token; decode: the single token).""" + cache_handle = engine_inputs.cache_manager + hidden = self._hidden(input_ids, position_ids, cache_handle) + logits = self.lm_head(hidden[-1:]) + return {"logits": [logits]} + + def can_batch(self, batch: NodeBatch, model_inputs: list[NodeInputs]) -> bool: + return True + + def forward_batched( + self, + graph_walk: str, + engine_inputs: ModelInputsFromEngine, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + **kwargs, + ) -> dict[str, NameToTensorList]: + """Batched forward: sample inside the pass (CUDA-graphable sampler) and + return per-request ``new_token``. Prefill packs all requests' tokens, so the + last-token-per-request indices come from the FlashInfer prefill wrapper's + persistent ``qo_indptr`` buffer; decode is one token per request.""" + cache_handle = engine_inputs.cache_manager + sampler = engine_inputs.sampler + cache_handle.set_active_label(_MAIN) + + hidden = self._hidden(input_ids, position_ids, cache_handle) + + if graph_walk == "prefill": + qo_indptr_buf = cache_handle.get_qo_indptr_buf(_MAIN) + assert qo_indptr_buf is not None, ( + "prefill forward_batched requires a CUDA-graph " + "FlashInferPrefillWrapper (qo_indptr static buffer); got None." + ) + last_token_indices = (qo_indptr_buf[1:] - 1).long() + hidden = hidden.index_select(0, last_token_indices) + elif graph_walk != "decode": + raise ValueError(f"Batched forward not supported for graph walk: {graph_walk!r}") + + logits = self.lm_head(hidden) # (bs, vocab) + request_ids = cache_handle.request_ids + new_tokens = self._sample(sampler, request_ids, logits) + return { + rid: {"new_token": [new_tokens[i : i + 1]]} + for i, rid in enumerate(request_ids) + } + + @staticmethod + def _sample( + sampler: Sampler, request_ids: list[str], logits: torch.Tensor, + ) -> torch.Tensor: + return sampler.sample(request_ids, logits, apply_penalty=True) + + def postprocess( + self, request_id: str, + request_info: CurrentForwardPassInfo, + outputs: dict[str, list[torch.Tensor]], + **kwargs, + ): + # Metadata-only: rebind the new token as the next step's text_inputs. + # EOS is checked in check_stop so the GPU thread doesn't sync on .item(). + if "new_token" not in outputs: + return + outputs["text_inputs"] = outputs["new_token"] + + def check_stop( + self, request_id: str, + request_info: CurrentForwardPassInfo, + outputs: dict[str, list[torch.Tensor]], + ) -> set[str]: + if "new_token" not in outputs: + return set() + token = outputs["new_token"][0].item() + eos_token_id = self.config.eos_token_id + ignore_eos = request_info.sampling_config["LLM"].ignore_eos + if (not ignore_eos and eos_token_id == token) or ( + request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1 + >= request_info.max_tokens + ): + return {"decode_loop"} + return set() diff --git a/mstar/model/kimi_k2_7/weight_loader.py b/mstar/model/kimi_k2_7/weight_loader.py new file mode 100644 index 000000000..d80d40537 --- /dev/null +++ b/mstar/model/kimi_k2_7/weight_loader.py @@ -0,0 +1,204 @@ +"""Kimi-K2.7 / DeepSeek-V3 weight loading (M5). + +Maps an HF ``DeepseekV3ForCausalLM`` checkpoint onto the mstar Kimi module tree +using the shared ``load_hf_weights`` machinery — a name remap plus stacked-shard +rules, exactly mirroring ``qwen3_omni_model.py::_get_thinker_stacked_params`` / +``_thinker_remap`` (same fused-expert layout). No shared abstraction is modified. + +Two transforms take the checkpoint keys to the module's ``named_parameters``: + +1. **Name remap** (:func:`kimi_name_remapper`): + - ``mlp.shared_experts.*`` -> ``mlp.shared_expert.*`` (HF plural -> our + singular ``ParallelGatedMLP`` submodule name); + - per-routed-expert ``mlp.experts.{i}.{gate,up,down}_proj.weight`` -> + ``mlp.experts.{gate,up,down}_proj.__expert{i}__.weight`` — the + ``__expert{i}__`` marker lets the stacked rules carry both the projection + *and* the expert slot in one ``shard_id``; + - everything else is identity (the checkpoint prefixes ``model.``, + ``self_attn.``, ``input_layernorm``, ``lm_head`` … already line up). + +2. **Stacked rules** (:func:`build_kimi_stacked_params`): + - routed experts: ``.experts.gate_proj.__expert{i}__.weight`` / + ``.up_proj.__expert{i}__.weight`` -> fused ``.experts.gate_up_proj`` + ("w13", gate then up) with ``shard_id="gate:i"/"up:i"``; + ``.down_proj.__expert{i}__.weight`` -> ``.experts.down_proj`` ("w2", + ``shard_id="down:i"``). The fused params get their per-shard + ``weight_loader`` in :class:`KimiSparseMoeBlock`. + - dense + shared SwiGLU ``.gate_proj`` / ``.up_proj`` -> merged + ``.gate_up_proj`` (shard 0 / 1). These MUST come *after* the expert rules: + ``_apply_stacked`` returns on first match and the remapped expert key + ``…experts.gate_proj.__expert{i}__.weight`` also contains ``.gate_proj``. + +**MLA loads strictly by name — NO q_a/kv_a fusion.** M3 built *separate* +``q_a_proj`` and ``kv_a_proj_with_mqa`` (the naive/materialized path), so their +checkpoint keys map straight to the identically-named params. Fusing them into a +single ``fused_qkv_a_proj`` is only needed for the deferred weight-absorbed MLA +class (``DeepseekV2MLAAttention``); the naive path needs no such fusion. This is +a deliberate simplification of the earlier plan note. + +**Router bias stays fp32.** ``KimiMoEGate.e_score_correction_bias`` is a +selection-only bias DeepSeek keeps in fp32 for router stability. A whole-model +``.to(bfloat16)`` would downcast it, so :func:`restore_router_bias_fp32` forces +every such param back to fp32 immediately before the load (the copy then lands +fp32 -> fp32). The checkpoint stores this tensor as fp32. + +The dense-vs-MoE split follows ``is_moe_layer`` (``first_k_dense_replace`` / +``moe_layer_freq``): early dense layers carry ``mlp.{gate,up,down}_proj`` into a +``ParallelGatedMLP``; MoE layers carry the expert-stacked + shared + gate params. +The routing here is layer-agnostic — a dense layer simply never emits +``mlp.experts.*`` / ``mlp.gate.*`` keys, and a MoE layer never emits a bare +``mlp.gate_proj``. + +Refs: HF key -> param authority is vLLM +``model_executor/models/deepseek_v2.py::DeepseekV2ForCausalLM.load_weights`` +(the ``stacked_params_mapping`` + per-expert ``expert_params_mapping`` there); +the fused ``w13``/``w2`` naming is vLLM's. + +---------------------------------------------------------------------------- +DESIGN NOTE — compressed-tensors INT4 / fp8 dequant-on-load (DEFERRED) +---------------------------------------------------------------------------- +The real Kimi-K2.7 checkpoint ships **compressed-tensors INT4** (fp8 variants +also exist); ``fused_experts`` is bf16/fp16-only and a full bf16 dequant of the +1T model is ~2 TB > 8xH200 (see the perf memo). So the real quantized checkpoint +is **out of scope for M5** (no checkpoint present, would not fit) — this loader is +the clean **bf16 path**, validated on a synthetic ``reduced()`` checkpoint. + +When the checkpoint + a quantized kernel land, dequant-on-load slots in **without +touching the routing above**, because ``load_hf_weights`` dispatches per parameter +through each param's ``weight_loader``: + + * compressed-tensors stores, per quantized tensor, a packed ``*.weight_packed`` + (INT4 nibbles / fp8 bytes) plus ``*.weight_scale`` (+ optional + ``*.weight_zero_point``) at a group granularity from the checkpoint's + ``quantization_config``. + * Hook point A (streaming): wrap the ``(name, tensor)`` iterator with a + dequantizer that consumes the ``weight_packed``/``weight_scale`` group, emits + a single bf16 ``*.weight`` (unpack nibble -> int -> ``(q - zp) * scale`` per + group), and drops the scale/zp/packed keys. Downstream routing is unchanged. + * Hook point B (per-param, memory-lean): keep the packed tensor in VRAM and + give the *destination* fused param a quant-aware ``weight_loader`` that stores + packed shards + scales (extend :class:`KimiSparseMoeBlock` to hold + ``gate_up_proj_packed`` / ``_scale``) and swap ``_dispatch`` for a quantized + grouped-GEMM. This is the only way to actually *serve* the 1T model and is + tracked as the top memory item in the perf backlog — a separate effort from + this DeepSeek-V3 port. + +Either hook is additive: the bf16 routing (remap + stacked rules) below is the +substrate both build on. +""" +from __future__ import annotations + +import re +from collections.abc import Iterable +from pathlib import Path + +import torch +from torch import nn + +from mstar.model.loader.base import StackedParamRule + +# HF checkpoint suffixes for the per-routed-expert projections. +_EXPERT_RE = re.compile( + r"(.*)\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$" +) + + +def kimi_name_remapper(name: str) -> str | None: + """HF DeepSeek-V3 checkpoint key -> Kimi module param path. + + Returns ``None`` to drop a key (precomputed ``rotary_emb`` buffers). See the + module docstring for the full mapping; MLA / norms / embed / lm_head are all + identity. + """ + if "rotary_emb" in name: + return None + # HF names the shared expert plural; our module has one ``shared_expert``. + name = name.replace(".shared_experts.", ".shared_expert.") + # Per-expert fusion marker so the stacked rules can pick up expert index. + m = _EXPERT_RE.match(name) + if m: + prefix, expert_idx, proj = m.groups() + return f"{prefix}.experts.{proj}.__expert{expert_idx}__.weight" + return name + + +def build_kimi_stacked_params(n_routed_experts: int) -> list[StackedParamRule]: + """Fused-shard routing for Kimi-K2.7 (mirrors the Qwen3-MoE thinker rules). + + Per-expert ``gate``/``up`` -> ``experts.gate_up_proj`` (w13) and ``down`` -> + ``experts.down_proj`` (w2), then the dense/shared SwiGLU gate/up merge. + Expert rules precede the dense rules (first-match wins in ``_apply_stacked``). + """ + rules: list[StackedParamRule] = [] + for i in range(n_routed_experts): + rules.append(StackedParamRule( + target_suffix=".experts.gate_up_proj", + source_suffix=f".experts.gate_proj.__expert{i}__.weight", + shard_id=f"gate:{i}", + )) + rules.append(StackedParamRule( + target_suffix=".experts.gate_up_proj", + source_suffix=f".experts.up_proj.__expert{i}__.weight", + shard_id=f"up:{i}", + )) + rules.append(StackedParamRule( + target_suffix=".experts.down_proj", + source_suffix=f".experts.down_proj.__expert{i}__.weight", + shard_id=f"down:{i}", + )) + # Dense MLP + shared-expert gate/up fusion — AFTER the expert rules. + rules.append(StackedParamRule(".gate_up_proj", ".gate_proj", 0)) + rules.append(StackedParamRule(".gate_up_proj", ".up_proj", 1)) + return rules + + +def restore_router_bias_fp32(module: nn.Module) -> None: + """Force every ``e_score_correction_bias`` back to fp32 in place. + + DeepSeek keeps this selection bias fp32; a whole-model ``.to(bfloat16)`` would + downcast it. Call immediately before loading so the source (fp32) copies into + an fp32 destination. + """ + for sub in module.modules(): + bias = getattr(sub, "e_score_correction_bias", None) + if isinstance(bias, nn.Parameter) and bias.dtype != torch.float32: + bias.data = bias.data.float() + + +def load_kimi_hf_weights( + module: nn.Module, + weights: Iterable[tuple[str, torch.Tensor]], + n_routed_experts: int, +) -> set[str]: + """Load an HF DeepSeek-V3 weight stream into ``module``. + + Thin wrapper: restore the fp32 router bias, then dispatch through + ``load_hf_weights`` with the Kimi remap + stacked rules. Returns the set of + param paths that received a tensor (callers can diff against + ``named_parameters()`` to assert completeness). + """ + from mstar.model.loader import load_hf_weights + + restore_router_bias_fp32(module) + return load_hf_weights( + module, + weights, + stacked_params=build_kimi_stacked_params(n_routed_experts), + name_remapper=kimi_name_remapper, + ) + + +def load_weights( + module: nn.Module, + source: str | Path, + device: torch.device | str = "cpu", +) -> set[str]: + """``(module, source, device)`` entrypoint mirroring Orpheus. + + ``source`` is a safetensors file or an HF-style checkpoint directory. Picks + the right streaming iterator and drives ``module.load_weights`` (which calls + :func:`load_kimi_hf_weights`). + """ + from mstar.model.loader import load_weights as _driver + + return _driver(module, source, device=device) diff --git a/mstar/model/registry.py b/mstar/model/registry.py index 87eedb2bf..8f16511b8 100644 --- a/mstar/model/registry.py +++ b/mstar/model/registry.py @@ -2,6 +2,7 @@ from mstar.model.base import Model from mstar.model.cosmos3.cosmos3_model import Cosmos3Model from mstar.model.higgs_audio.higgs_audio_model import HiggsAudioModel +from mstar.model.kimi_k2_7.kimi_model import KimiK2Model from mstar.model.orpheus.orpheus_model import OrpheusModel from mstar.model.pi05.pi05_model import Pi05Model from mstar.model.qwen3_omni.qwen3_omni_model import Qwen3OmniModel @@ -16,6 +17,7 @@ "cosmos3_droid": Cosmos3Model, "cosmos3_super": Cosmos3Model, "higgs_audio": HiggsAudioModel, + "kimi_k2_7": KimiK2Model, "orpheus": OrpheusModel, "pi05": Pi05Model, "qwen3_omni": Qwen3OmniModel, @@ -42,6 +44,9 @@ # Higgs-Audio v3 STT: Whisper-style audio tower + Qwen3-1.7B LLM. # (The v2 checkpoints are TTS/generation models, not ASR.) "higgs_audio": {"model_path_hf": "bosonai/higgs-audio-v3-stt"}, + # Kimi-K2.7-Code: 1T MoE (DeepSeek-V3 text backbone + MoonViT). M0 is a + # text-only scaffold; real serving is TP8 / multi-node. + "kimi_k2_7": {"model_path_hf": "moonshotai/Kimi-K2.7-Code"}, "orpheus": {"model_path_hf": "canopylabs/orpheus-3b-0.1-ft"}, # Pi0.5 PyTorch port published by lerobot — single safetensors blob # (~14 GB). mstar/model/pi05/weight_loader.py handles the lerobot->mstar diff --git a/test/integration/test_kimi_components.py b/test/integration/test_kimi_components.py new file mode 100644 index 000000000..8cf9d6a3a --- /dev/null +++ b/test/integration/test_kimi_components.py @@ -0,0 +1,148 @@ +"""M1 golden tests for Kimi-K2.7 cheap reused components. + +For each cheap component of the DeepSeek-V3 text backbone — RMSNorm, the dense +SwiGLU MLP, the token embedding, and the LM head — build the mstar component from +``KimiK2Config.reduced()``, load identical random weights into it and an +independent reference, and assert the outputs match. The reference formulas are +inlined here (self-contained; no dependency on the local golden harness) and each +is cited to the vLLM DeepSeek-V3 source so the golden is authoritative. + +This is a GPU test: mstar's standard RMSNorm dispatches to a FlashInfer fused +kernel, so the suite runs on ``cuda``. It skips automatically without a GPU. + +Run: pytest test/integration/test_kimi_components.py -v +""" +import pytest +import torch +import torch.nn.functional as F + +from mstar.model.kimi_k2_7.components.language_model import ( + build_dense_mlp, + build_embedding, + build_lm_head, + build_rmsnorm, +) +from mstar.model.kimi_k2_7.config import KimiK2Config + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="M1 golden tests need a GPU (mstar RMSNorm uses a FlashInfer kernel)", +) + +DEVICE = "cuda" + + +def _cfg() -> KimiK2Config: + return KimiK2Config.reduced() + + +# -------------------------------------------------------------------------- +# Independent references (cited to vllm-project/vllm .../models/deepseek_v2.py) +# -------------------------------------------------------------------------- + +def _ref_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + # Standard Llama/DeepSeek RMSNorm (HF ``LlamaRMSNorm`` / vLLM + # ``RMSNorm.forward_native``): normalize in fp32, scale by weight in the + # input dtype. + orig_dtype = x.dtype + x32 = x.float() + var = x32.pow(2).mean(dim=-1, keepdim=True) + x32 = x32 * torch.rsqrt(var + eps) + return weight * x32.to(orig_dtype) + + +def _ref_swiglu( + x: torch.Tensor, gate_w: torch.Tensor, up_w: torch.Tensor, down_w: torch.Tensor +) -> torch.Tensor: + # ``DeepseekV2MLP.forward`` = ``down_proj(SiluAndMul(gate_up_proj(x)))``, + # bias=False, silu-only. ``SiluAndMul([g, u]) = silu(g) * u``. + gate = F.linear(x, gate_w) + up = F.linear(x, up_w) + return F.linear(F.silu(gate) * up, down_w) + + +def _ref_embedding(ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return F.embedding(ids, weight) + + +def _ref_lm_head(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + # Untied head (``tie_word_embeddings=False``): plain ``x @ weight.T``. + return F.linear(x, weight) + + +# -------------------------------------------------------------------------- +# Tests +# -------------------------------------------------------------------------- + +def test_rmsnorm_matches_reference(): + torch.manual_seed(0) + cfg = _cfg() + dtype = torch.bfloat16 # FlashInfer rmsnorm runs in half precision + n_tokens = 7 + + norm = build_rmsnorm(cfg).to(device=DEVICE, dtype=dtype) + weight = torch.randn(cfg.hidden_size, device=DEVICE, dtype=dtype) + norm.weight.data.copy_(weight) + + x = torch.randn(n_tokens, cfg.hidden_size, device=DEVICE, dtype=dtype) + + got = norm(x) + expected = _ref_rmsnorm(x, weight, cfg.rms_norm_eps) + torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2) + + +def test_dense_mlp_matches_reference(): + torch.manual_seed(1) + cfg = _cfg() + dtype = torch.float32 + n_tokens = 5 + h, i = cfg.hidden_size, cfg.intermediate_size + + mlp = build_dense_mlp(cfg).to(device=DEVICE, dtype=dtype) + gate_w = torch.randn(i, h, device=DEVICE, dtype=dtype) * 0.05 + up_w = torch.randn(i, h, device=DEVICE, dtype=dtype) * 0.05 + down_w = torch.randn(h, i, device=DEVICE, dtype=dtype) * 0.05 + # Load through the real (stacked) loaders: gate -> shard 0, up -> shard 1. + mlp.gate_up_proj.weight_loader(mlp.gate_up_proj.weight, gate_w, loaded_shard_id=0) + mlp.gate_up_proj.weight_loader(mlp.gate_up_proj.weight, up_w, loaded_shard_id=1) + mlp.down_proj.weight_loader(mlp.down_proj.weight, down_w) + + x = torch.randn(n_tokens, h, device=DEVICE, dtype=dtype) + + got = mlp(x) + expected = _ref_swiglu(x, gate_w, up_w, down_w) + torch.testing.assert_close(got, expected, rtol=1e-4, atol=1e-4) + + +def test_embedding_matches_reference(): + torch.manual_seed(2) + cfg = _cfg() + dtype = torch.float32 + + emb = build_embedding(cfg).to(device=DEVICE, dtype=dtype) + weight = torch.randn(cfg.vocab_size, cfg.hidden_size, device=DEVICE, dtype=dtype) + emb.weight_loader(emb.weight, weight) + + ids = torch.randint(0, cfg.vocab_size, (9,), device=DEVICE) + + got = emb(ids) + expected = _ref_embedding(ids, weight) + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + +def test_lm_head_matches_reference(): + torch.manual_seed(3) + cfg = _cfg() + dtype = torch.float32 + n_tokens = 4 + + head = build_lm_head(cfg).to(device=DEVICE, dtype=dtype) + weight = torch.randn(cfg.vocab_size, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.02 + head.weight_loader(head.weight, weight) + + x = torch.randn(n_tokens, cfg.hidden_size, device=DEVICE, dtype=dtype) + + got = head(x) + expected = _ref_lm_head(x, weight) + assert got.shape == (n_tokens, cfg.vocab_size) + torch.testing.assert_close(got, expected, rtol=1e-4, atol=1e-4) diff --git a/test/integration/test_kimi_decoder_layer.py b/test/integration/test_kimi_decoder_layer.py new file mode 100644 index 000000000..d4b5598cf --- /dev/null +++ b/test/integration/test_kimi_decoder_layer.py @@ -0,0 +1,261 @@ +"""M4 golden tests for the Kimi-K2.7 / DeepSeek-V3 decoder layer. + +One golden per feed-forward variant — a dense layer (``layer_idx=0``, below +``first_k_dense_replace``) and a MoE layer (``layer_idx=1``) — each compared to a +self-contained inline reference that re-derives the whole block: +pre-norm → naive-MLA self-attention → residual → pre-norm → (dense-or-MoE) FFN → +residual. The inner attention and FFN references are the same ones the M2/M3 +goldens use, cited to vLLM; what this test adds is the residual/norm *wiring*, +matching vLLM ``models/deepseek_v2.py::DeepseekV2DecoderLayer.forward``. + +A ``_MockMLACache`` stands in for the paged cache (causal SDPA at the fixed +``1/sqrt(qk_head_dim)`` scale FlashInfer uses); the real paged ``run_attention`` +is exercised separately in ``test_kimi_flashinfer_attention.py``. + +GPU test (mstar RMSNorm + the fused expert GEMM are CUDA/half-precision only); +skips without a GPU. + +Run: pytest test/integration/test_kimi_decoder_layer.py -v +""" +import pytest +import torch +import torch.nn.functional as F + +from mstar.model.kimi_k2_7.components.decoder_layer import KimiDecoderLayer +from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock +from mstar.model.kimi_k2_7.components.rope import ( + _yarn_find_correction_range, + _yarn_linear_ramp_mask, + rotate_gptj, + yarn_get_mscale, +) +from mstar.model.kimi_k2_7.config import KimiK2Config + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="M4 golden tests need a GPU (RMSNorm + fused expert GEMM are CUDA-only)", +) + +DEVICE = "cuda" + + +# -------------------------------------------------------------------------- +# Inline references (cited to vLLM deepseek_v2.py / deepseek_scaling_rope.py / +# cpu_fused_moe.py) — self-contained, no dependency on the golden harness. +# -------------------------------------------------------------------------- + +def _ref_rmsnorm(x, weight, eps): + x32 = x.float() + x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) + return weight * x32.to(x.dtype) + + +def _ref_yarn_rope(pos, q_pe, k_pe, rotary_dim, base, factor, max_pos, + beta_fast, beta_slow, mscale, mscale_all_dim): + pos_freqs = base ** (torch.arange(0, rotary_dim, 2, device=q_pe.device).float() / rotary_dim) + ext, interp = 1.0 / pos_freqs, 1.0 / (factor * pos_freqs) + low, high = _yarn_find_correction_range(beta_fast, beta_slow, rotary_dim, base, max_pos) + mask = 1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float).to(q_pe.device) + inv_freq = interp * (1 - mask) + ext * mask + amp = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim) + freqs = torch.outer(pos.float(), inv_freq) + cos = (freqs.cos() * amp).repeat_interleave(2, -1).unsqueeze(-2) + sin = (freqs.sin() * amp).repeat_interleave(2, -1).unsqueeze(-2) + qr = q_pe.float() * cos + rotate_gptj(q_pe.float()) * sin + kr = k_pe.float() * cos + rotate_gptj(k_pe.float()) * sin + return qr.to(q_pe.dtype), kr.to(k_pe.dtype) + + +def _sdpa_causal(q, k, v, scale): + """Causal SDPA at a fixed scale (mirrors _MockMLACache / FlashInfer).""" + qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) # (H,T,D) + T = q.shape[0] + causal = torch.triu( + torch.full((T, T), float("-inf"), device=q.device), diagonal=1) + attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1) + return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) + + +def _ref_attn_forward(attn, cfg, h_normed, pos): + """Independent naive-MLA forward matching KimiMLAAttention.forward.""" + T, H = h_normed.shape[0], attn.num_heads + Dnope, Drope, Dv, L = ( + cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank) + eps = cfg.rms_norm_eps + q = _ref_rmsnorm(F.linear(h_normed, attn.q_a_proj.weight), attn.q_a_layernorm.weight, eps) + q = F.linear(q, attn.q_b_proj.weight).view(T, H, cfg.qk_head_dim) + q_nope, q_pe = q.split([Dnope, Drope], dim=-1) + latent = F.linear(h_normed, attn.kv_a_proj_with_mqa.weight) + kv_a, k_pe = latent.split([L, Drope], dim=-1) + kv = F.linear(_ref_rmsnorm(kv_a, attn.kv_a_layernorm.weight, eps), + attn.kv_b_proj.weight).view(T, H, Dnope + Dv) + k_nope, v = kv.split([Dnope, Dv], dim=-1) + k_pe = k_pe.view(T, 1, Drope) + r = cfg.rope_scaling + q_pe, k_pe = _ref_yarn_rope( + pos, q_pe, k_pe, Drope, cfg.rope_theta, r["factor"], + r["original_max_position_embeddings"], r.get("beta_fast", 32), + r.get("beta_slow", 1), r.get("mscale", 1.0), r.get("mscale_all_dim", 0.0)) + q = torch.cat([q_nope, q_pe], dim=-1) * attn.softmax_scale_boost + k = torch.cat([k_nope, k_pe.expand(T, H, Drope)], dim=-1) + v = F.pad(v, [0, cfg.qk_head_dim - Dv]) + out = _sdpa_causal(q, k, v, cfg.qk_head_dim ** -0.5) + out = out[..., :Dv].reshape(T, H * Dv) + return F.linear(out, attn.o_proj.weight) + + +def _ref_grouped_topk(logits, bias, n_group, topk_group, top_k, + norm_topk_prob, routed_scaling_factor): + scores = logits.float().sigmoid() + T = scores.shape[0] + original = scores + scores = scores + bias.unsqueeze(0) + group_scores = scores.view(T, n_group, -1).topk(2, dim=-1)[0].sum(dim=-1) + group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = ( + group_mask.unsqueeze(-1) + .expand(T, n_group, scores.shape[-1] // n_group) + .reshape(T, -1) + ) + masked = scores.masked_fill(~score_mask.bool(), float("-inf")) + ids = torch.topk(masked, k=top_k, dim=-1, sorted=False)[1] + weights = original.gather(1, ids) + if norm_topk_prob: + weights = weights / weights.sum(dim=-1, keepdim=True) + weights = weights * routed_scaling_factor + return weights, ids + + +def _ref_routed_experts(h, gate_up, down, weights, ids): + T, H = h.shape + inter = down.shape[-1] + out = torch.zeros(T, H, dtype=h.dtype, device=h.device) + for t in range(T): + for j in range(ids.shape[1]): + e = int(ids[t, j]) + gu = gate_up[e] @ h[t] + g, u = gu[:inter], gu[inter:] + out[t] += weights[t, j] * (down[e] @ (F.silu(g) * u)) + return out + + +def _ref_swiglu(x, gate_w, up_w, down_w): + return F.linear(F.silu(F.linear(x, gate_w)) * F.linear(x, up_w), down_w) + + +def _ref_mlp_forward(mlp, cfg, h_normed): + """Dense SwiGLU or MoE (routed + ungated shared), matching the module.""" + if isinstance(mlp, KimiSparseMoeBlock): + I = cfg.moe_intermediate_size + si = cfg.moe_intermediate_size * cfg.n_shared_experts + logits = F.linear(h_normed.float(), mlp.gate.weight.float()) + weights, ids = _ref_grouped_topk( + logits, mlp.gate.e_score_correction_bias, cfg.n_group, cfg.topk_group, + cfg.num_experts_per_tok, cfg.norm_topk_prob, cfg.routed_scaling_factor) + routed = _ref_routed_experts( + h_normed, mlp.experts.gate_up_proj, mlp.experts.down_proj, + weights.to(h_normed.dtype), ids) + sh = mlp.shared_expert + shared = _ref_swiglu( + h_normed, sh.gate_up_proj.weight[:si], sh.gate_up_proj.weight[si:], + sh.down_proj.weight) + return routed + shared + i = cfg.intermediate_size + return _ref_swiglu( + h_normed, mlp.gate_up_proj.weight[:i], mlp.gate_up_proj.weight[i:], + mlp.down_proj.weight) + + +def _ref_decoder_layer(layer, cfg, h, pos): + eps = cfg.rms_norm_eps + attn_in = _ref_rmsnorm(h, layer.input_layernorm.weight, eps) + h1 = h + _ref_attn_forward(layer.self_attn, cfg, attn_in, pos) + mlp_in = _ref_rmsnorm(h1, layer.post_attention_layernorm.weight, eps) + return h1 + _ref_mlp_forward(layer.mlp, cfg, mlp_in) + + +# -------------------------------------------------------------------------- +# Mock paged cache: causal SDPA at 1/sqrt(head_dim), no cross-layer history +# (a single prefill forward; each layer attends its own q/k/v). +# -------------------------------------------------------------------------- + +class _MockMLACache: + def __init__(self, head_dim): + self.scale = head_dim ** -0.5 + + def set_layer_idx(self, _i): + pass + + def advance_seq_lens(self, *_a, **_k): + pass + + def run_attention(self, q, k, v): + return _sdpa_causal(q, k, v, self.scale) + + +def _build_layer(cfg, layer_idx, dtype): + layer = KimiDecoderLayer(cfg, layer_idx).to(device=DEVICE, dtype=dtype) + a = layer.self_attn + for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): + lin.weight.data.normal_(0, 0.03) + for norm in (a.q_a_layernorm, a.kv_a_layernorm): + norm.weight.data.normal_(1.0, 0.02) + layer.input_layernorm.weight.data.normal_(1.0, 0.02) + layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02) + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + # Keep the router fp32 (deterministic selection); experts/shared bf16. + mlp.gate.weight.data = torch.randn( + cfg.n_routed_experts, cfg.hidden_size, device=DEVICE) + mlp.gate.e_score_correction_bias.data = torch.randn( + cfg.n_routed_experts, device=DEVICE) + mlp.experts.gate_up_proj.data.normal_(0, 0.05) + mlp.experts.down_proj.data.normal_(0, 0.05) + mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05) + else: + mlp.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.down_proj.weight.data.normal_(0, 0.05) + return layer + + +# -------------------------------------------------------------------------- +# Tests +# -------------------------------------------------------------------------- + +def test_dense_decoder_layer_matches_reference(): + torch.manual_seed(0) + cfg = KimiK2Config.reduced() + dtype = torch.bfloat16 + layer = _build_layer(cfg, layer_idx=0, dtype=dtype) # dense (< first_k_dense_replace) + assert not isinstance(layer.mlp, KimiSparseMoeBlock) + + T = 6 + h = torch.randn(T, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.1 + pos = torch.arange(T, device=DEVICE) + + got = layer(h, _MockMLACache(cfg.qk_head_dim), pos) + expected = _ref_decoder_layer(layer, cfg, h, pos) + + assert got.shape == (T, cfg.hidden_size) + torch.testing.assert_close(got, expected, rtol=3e-2, atol=3e-2) + + +def test_moe_decoder_layer_matches_reference(): + torch.manual_seed(1) + cfg = KimiK2Config.reduced() + dtype = torch.bfloat16 + layer = _build_layer(cfg, layer_idx=1, dtype=dtype) # MoE (>= first_k_dense_replace) + assert isinstance(layer.mlp, KimiSparseMoeBlock) + + T = 7 + h = torch.randn(T, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.1 + pos = torch.arange(T, device=DEVICE) + + got = layer(h, _MockMLACache(cfg.qk_head_dim), pos) + expected = _ref_decoder_layer(layer, cfg, h, pos) + + assert got.shape == (T, cfg.hidden_size) + torch.testing.assert_close(got, expected, rtol=3e-2, atol=3e-2) diff --git a/test/integration/test_kimi_flashinfer_attention.py b/test/integration/test_kimi_flashinfer_attention.py new file mode 100644 index 000000000..c4bb9180c --- /dev/null +++ b/test/integration/test_kimi_flashinfer_attention.py @@ -0,0 +1,169 @@ +"""M4 step 3: validate mstar's REAL paged ``run_attention`` for the naive MLA. + +The naive/materialized MLA stores a ``head_dim = qk_head_dim`` (nope+rope) K plus +a V padded to that same width, then calls the paged ``run_attention`` at the +fixed ``1/sqrt(head_dim)`` scale FlashInfer uses (which is exactly why the +``mscale^2`` softmax boost is folded into q — ``run_attention`` exposes no custom +``sm_scale``). This test drives the **real** ``FlashInferCacheManager`` over a +genuine paged KV cache (real ``PagedAllocationManager`` + ``LocalTransferEngine``) +and asserts its ``run_attention`` matches a causal-SDPA reference at +``1/sqrt(head_dim)`` — confirming both the paged path integrates and the scale +assumption the naive MLA relies on. + +KEY CONSTRAINT FOUND (this is the "FlashInfer-192" answer, recorded for M5/M6): +FlashInfer 0.6.14's SM90 (Hopper / H200) prefill kernel has a compile-time +``static_assert(HEAD_DIM_VO == 64 || HEAD_DIM_VO == 128 || HEAD_DIM_VO == 256)`` +(``flashinfer/.../attention/hopper/prefill_sm90.cuh:572``). The naive MLA pads V +to ``qk_head_dim``, so ``head_dim_vo == head_dim``: + + * real Kimi ``qk_head_dim = 192`` (nope 128 + rope 64) -> vo=192 -> JIT FAILS + * reduced-config ``qk_head_dim = 24`` -> vo=24 -> JIT FAILS + * 64 / 128 / 256 -> supported -> OK + +So the naive MLA path cannot use the paged ``run_attention`` at head_dim 192 (or +the reduced 24) on Hopper as-is. Validated mitigation (M6 follow-up, NOT done +here): pad ``head_dim`` up to the next supported vo (256 for real 192, 64 for the +reduced 24), pad q/k/v to it, and slice the attention output back to +``v_head_dim``. The supported-dim runs below (128 / 256) are exactly that padded +path; the env-gated test at the bottom records the raw 192 failure. + +Run: pytest test/integration/test_kimi_flashinfer_attention.py -v + KIMI_TEST_FLASHINFER_192=1 pytest ... -k rejects # ~60s failing JIT +""" +import os + +import pytest +import torch + +from mstar.communication.tensors import LocalTransferEngine +from mstar.engine.cache_manager import WorkspaceBufferManager, create_cache_manager +from mstar.engine.kv_store import ( + KVCacheConfig, + PagedAllocationManager, + TransferEngineInfo, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="real FlashInfer paged attention needs a GPU", +) + +DEVICE = torch.device("cuda") + + +def _make_real_cache_manager(num_heads, head_dim, dtype, page_size=128, max_num_pages=8): + """Build a genuine paged FlashInferCacheManager for one request. + + Mirrors ``KVCacheEngine.load_model`` / ``_create_cache_manager``: a real + ``[layers, pages, 2, page_size, heads, head_dim]`` KV cache, a + ``PagedAllocationManager`` over a no-op ``LocalTransferEngine`` (single-node + SHM path — no cross-worker reads), and the flashinfer backend. Returns + ``(cache_manager, alloc_manager)`` so the caller can clean up. + """ + kv_cache = torch.zeros( + 2, max_num_pages, 2, page_size, num_heads, head_dim, + dtype=dtype, device=DEVICE, + ).contiguous() + kv_cfg = KVCacheConfig( + num_layers=2, num_kv_heads=num_heads, head_dim=head_dim, + max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages, + page_size=page_size, num_qo_heads=num_heads, + ) + transfer_info = TransferEngineInfo( + my_entity_id="kimi_flashinfer_test", + my_session_id="kimi_session", + transfer_engine=LocalTransferEngine("localhost"), + ) + alloc = PagedAllocationManager( + config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info, + ) + alloc.add_request("r0", ["main"]) + buffers = WorkspaceBufferManager(64 * 1024 * 1024, device=DEVICE) + cm = create_cache_manager( + request_ids=["r0"], + active_labels_per_request={"r0": "main"}, + kv_cache=kv_cache, + alloc_manager=alloc, + buffer_manager=buffers, + kv_cache_config=kv_cfg, + device=DEVICE, + ) + return cm, alloc + + +def _sdpa_causal(q, k, v, scale): + qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) # (H,T,D) + T = q.shape[0] + causal = torch.triu( + torch.full((T, T), float("-inf"), device=q.device), diagonal=1) + attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1) + return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) + + +# FlashInfer SM90 prefill supports head_dim_vo in {64, 128, 256}. These are the +# sizes the naive-MLA V-pad would target: 64 for the reduced config (24 -> 64), +# 128 canonical, 256 for the real Kimi qk_head_dim (192 -> 256). +@pytest.mark.parametrize("head_dim", [128, 256]) +def test_real_paged_run_attention_matches_sdpa(head_dim): + """The real FlashInferCacheManager.run_attention == causal SDPA at + 1/sqrt(head_dim), the fixed scale the naive MLA folds mscale^2 into q for.""" + torch.manual_seed(0) + num_heads, T = 4, 6 + dtype = torch.bfloat16 + cm, alloc = _make_real_cache_manager(num_heads, head_dim, dtype) + try: + q = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1 + k = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1 + v = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1 + + cm.set_active_label("main") + cm.plan_attention(seq_lens=[T], is_causal=True, dtype=dtype) + cm.set_layer_idx(0) + got = cm.run_attention(q=q, k=k, v=v) + torch.cuda.synchronize() + + expected = _sdpa_causal(q, k, v, head_dim ** -0.5) + assert got.shape == (T, num_heads, head_dim) + torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2) + finally: + alloc.cleanup() + + +@pytest.mark.skipif( + os.environ.get("KIMI_TEST_FLASHINFER_192") != "1", + reason="opt-in (~60s failing JIT): set KIMI_TEST_FLASHINFER_192=1 to record " + "the head_dim=192 SM90 static_assert rejection", +) +def test_flashinfer_rejects_head_dim_192(): + """Executable record of the constraint: the real paged run_attention cannot + JIT-build for head_dim=192 (vo=192) on Hopper — FlashInfer static_asserts + HEAD_DIM_VO in {64,128,256}. If FlashInfer/the mitigation ever lifts this, + this test flips to failing (no exception raised) and flags the change. + + The failure surfaces as the JIT build erroring out; the concrete exception + type varies by stage (a RuntimeError wrapping the ninja/nvcc + CalledProcessError), so we assert on the broad base and check the message + points at the build rather than an unrelated error.""" + torch.manual_seed(0) + num_heads, T, head_dim = 4, 6, 192 + dtype = torch.bfloat16 + cm, alloc = _make_real_cache_manager(num_heads, head_dim, dtype) + try: + q = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1 + k = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1 + v = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1 + cm.set_active_label("main") + # The offending kernel is JIT-built when FlashInfer schedules it — that + # can happen in plan_attention (the wrapper.plan() call) or run_attention + # depending on version, so both sit inside the raises block. + with pytest.raises(Exception) as exc_info: + cm.plan_attention(seq_lens=[T], is_causal=True, dtype=dtype) + cm.set_layer_idx(0) + cm.run_attention(q=q, k=k, v=v) + torch.cuda.synchronize() + # Guard against catching an unrelated error: the message must reference + # the failed build / the offending head_dim. + msg = str(exc_info.value).lower() + assert "ninja" in msg or "build" in msg or "192" in msg + finally: + alloc.cleanup() diff --git a/test/integration/test_kimi_forward.py b/test/integration/test_kimi_forward.py new file mode 100644 index 000000000..ce426b940 --- /dev/null +++ b/test/integration/test_kimi_forward.py @@ -0,0 +1,271 @@ +"""M4 full-forward golden test for Kimi-K2.7 / DeepSeek-V3 (assembled backbone). + +Runs ``KimiForCausalLM`` end to end on the reduced config — token ids → +embedding → stacked ``KimiDecoderLayer`` blocks (including the dense→MoE +transition at ``first_k_dense_replace=1``) → final RMSNorm → untied LM head → +logits — and compares against a self-contained inline reference that re-derives +every step. The inner attention / FFN / router references are the same ones the +M2/M3 goldens use (cited to vLLM); this test verifies the *assembly*: embedding, +the per-layer cache-handle contract (``set_layer_idx`` each layer, +``advance_seq_lens`` once after), the layer stack, and the LM head. + +A ``_MockMLACache`` stands in for the paged cache (causal SDPA at the fixed +``1/sqrt(qk_head_dim)`` scale). Because a single prefill forward re-attends the +same tokens each layer with no cross-layer history, the mock — which attends the +q/k/v of each ``run_attention`` call independently — reproduces the paged +prefill exactly. The real FlashInfer paged path is validated separately in +``test_kimi_flashinfer_attention.py``. + +Refs: vLLM ``models/deepseek_v2.py`` (``DeepseekV2Model`` / ``DecoderLayer`` / +``DeepseekV2MoE``), ``rotary_embedding/deepseek_scaling_rope.py``, +``fused_moe/cpu_fused_moe.py::grouped_topk``. + +Run: pytest test/integration/test_kimi_forward.py -v +""" +import pytest +import torch +import torch.nn.functional as F + +from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM +from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock +from mstar.model.kimi_k2_7.components.rope import ( + _yarn_find_correction_range, + _yarn_linear_ramp_mask, + rotate_gptj, + yarn_get_mscale, +) +from mstar.model.kimi_k2_7.config import KimiK2Config + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="M4 full-forward golden needs a GPU (RMSNorm + fused expert GEMM)", +) + +DEVICE = "cuda" + + +# -------------------------------------------------------------------------- +# Inline references (self-contained; cited to vLLM). Same math as the M2/M3 +# component goldens, assembled here into a whole-model forward. +# -------------------------------------------------------------------------- + +def _ref_rmsnorm(x, weight, eps): + x32 = x.float() + x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) + return weight * x32.to(x.dtype) + + +def _ref_yarn_rope(pos, q_pe, k_pe, rotary_dim, base, factor, max_pos, + beta_fast, beta_slow, mscale, mscale_all_dim): + pos_freqs = base ** (torch.arange(0, rotary_dim, 2, device=q_pe.device).float() / rotary_dim) + ext, interp = 1.0 / pos_freqs, 1.0 / (factor * pos_freqs) + low, high = _yarn_find_correction_range(beta_fast, beta_slow, rotary_dim, base, max_pos) + mask = 1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float).to(q_pe.device) + inv_freq = interp * (1 - mask) + ext * mask + amp = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim) + freqs = torch.outer(pos.float(), inv_freq) + cos = (freqs.cos() * amp).repeat_interleave(2, -1).unsqueeze(-2) + sin = (freqs.sin() * amp).repeat_interleave(2, -1).unsqueeze(-2) + qr = q_pe.float() * cos + rotate_gptj(q_pe.float()) * sin + kr = k_pe.float() * cos + rotate_gptj(k_pe.float()) * sin + return qr.to(q_pe.dtype), kr.to(k_pe.dtype) + + +def _sdpa_causal(q, k, v, scale): + qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) + T = q.shape[0] + causal = torch.triu( + torch.full((T, T), float("-inf"), device=q.device), diagonal=1) + attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1) + return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) + + +def _ref_attn_forward(attn, cfg, h_normed, pos): + T, H = h_normed.shape[0], attn.num_heads + Dnope, Drope, Dv, L = ( + cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank) + eps = cfg.rms_norm_eps + q = _ref_rmsnorm(F.linear(h_normed, attn.q_a_proj.weight), attn.q_a_layernorm.weight, eps) + q = F.linear(q, attn.q_b_proj.weight).view(T, H, cfg.qk_head_dim) + q_nope, q_pe = q.split([Dnope, Drope], dim=-1) + latent = F.linear(h_normed, attn.kv_a_proj_with_mqa.weight) + kv_a, k_pe = latent.split([L, Drope], dim=-1) + kv = F.linear(_ref_rmsnorm(kv_a, attn.kv_a_layernorm.weight, eps), + attn.kv_b_proj.weight).view(T, H, Dnope + Dv) + k_nope, v = kv.split([Dnope, Dv], dim=-1) + k_pe = k_pe.view(T, 1, Drope) + r = cfg.rope_scaling + q_pe, k_pe = _ref_yarn_rope( + pos, q_pe, k_pe, Drope, cfg.rope_theta, r["factor"], + r["original_max_position_embeddings"], r.get("beta_fast", 32), + r.get("beta_slow", 1), r.get("mscale", 1.0), r.get("mscale_all_dim", 0.0)) + # M6 mitigation: q/k padded from Dqk and v from Dv up to padded_head_dim; the + # softmax_scale_boost compensates so run_attention's padded_head_dim**-0.5 + # scale reproduces the DeepSeek qk_head_dim**-0.5 * mscale**2 scale. + pad = cfg.padded_head_dim + q = F.pad(torch.cat([q_nope, q_pe], dim=-1), [0, pad - cfg.qk_head_dim]) * attn.softmax_scale_boost + k = F.pad(torch.cat([k_nope, k_pe.expand(T, H, Drope)], dim=-1), [0, pad - cfg.qk_head_dim]) + v = F.pad(v, [0, pad - Dv]) + out = _sdpa_causal(q, k, v, cfg.padded_head_dim ** -0.5) + out = out[..., :Dv].reshape(T, H * Dv) + return F.linear(out, attn.o_proj.weight) + + +def _ref_grouped_topk(logits, bias, n_group, topk_group, top_k, + norm_topk_prob, routed_scaling_factor): + scores = logits.float().sigmoid() + T = scores.shape[0] + original = scores + scores = scores + bias.unsqueeze(0) + group_scores = scores.view(T, n_group, -1).topk(2, dim=-1)[0].sum(dim=-1) + group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = ( + group_mask.unsqueeze(-1) + .expand(T, n_group, scores.shape[-1] // n_group) + .reshape(T, -1) + ) + masked = scores.masked_fill(~score_mask.bool(), float("-inf")) + ids = torch.topk(masked, k=top_k, dim=-1, sorted=False)[1] + weights = original.gather(1, ids) + if norm_topk_prob: + weights = weights / weights.sum(dim=-1, keepdim=True) + weights = weights * routed_scaling_factor + return weights, ids + + +def _ref_routed_experts(h, gate_up, down, weights, ids): + T, H = h.shape + inter = down.shape[-1] + out = torch.zeros(T, H, dtype=h.dtype, device=h.device) + for t in range(T): + for j in range(ids.shape[1]): + e = int(ids[t, j]) + gu = gate_up[e] @ h[t] + g, u = gu[:inter], gu[inter:] + out[t] += weights[t, j] * (down[e] @ (F.silu(g) * u)) + return out + + +def _ref_swiglu(x, gate_w, up_w, down_w): + return F.linear(F.silu(F.linear(x, gate_w)) * F.linear(x, up_w), down_w) + + +def _ref_mlp_forward(mlp, cfg, h_normed): + if isinstance(mlp, KimiSparseMoeBlock): + si = cfg.moe_intermediate_size * cfg.n_shared_experts + logits = F.linear(h_normed.float(), mlp.gate.weight.float()) + weights, ids = _ref_grouped_topk( + logits, mlp.gate.e_score_correction_bias, cfg.n_group, cfg.topk_group, + cfg.num_experts_per_tok, cfg.norm_topk_prob, cfg.routed_scaling_factor) + routed = _ref_routed_experts( + h_normed, mlp.experts.gate_up_proj, mlp.experts.down_proj, + weights.to(h_normed.dtype), ids) + sh = mlp.shared_expert + shared = _ref_swiglu( + h_normed, sh.gate_up_proj.weight[:si], sh.gate_up_proj.weight[si:], + sh.down_proj.weight) + return routed + shared + i = cfg.intermediate_size + return _ref_swiglu( + h_normed, mlp.gate_up_proj.weight[:i], mlp.gate_up_proj.weight[i:], + mlp.down_proj.weight) + + +def _ref_decoder_layer(layer, cfg, h, pos): + eps = cfg.rms_norm_eps + attn_in = _ref_rmsnorm(h, layer.input_layernorm.weight, eps) + h1 = h + _ref_attn_forward(layer.self_attn, cfg, attn_in, pos) + mlp_in = _ref_rmsnorm(h1, layer.post_attention_layernorm.weight, eps) + return h1 + _ref_mlp_forward(layer.mlp, cfg, mlp_in) + + +def _ref_forward(model, cfg, ids, pos): + h = F.embedding(ids, model.model.embed_tokens.weight) + for layer in model.model.layers: + h = _ref_decoder_layer(layer, cfg, h, pos) + h = _ref_rmsnorm(h, model.model.norm.weight, cfg.rms_norm_eps) + return F.linear(h, model.lm_head.weight) + + +# -------------------------------------------------------------------------- +# Mock paged cache (causal SDPA at 1/sqrt(head_dim)) + weight init +# -------------------------------------------------------------------------- + +class _MockMLACache: + def __init__(self, head_dim): + self.scale = head_dim ** -0.5 + self.layer_idx = 0 + self.advance_calls = 0 + + def set_layer_idx(self, i): + self.layer_idx = i + + def advance_seq_lens(self, *_a, **_k): + self.advance_calls += 1 + + def run_attention(self, q, k, v): + return _sdpa_causal(q, k, v, self.scale) + + +def _fill_layer(layer, cfg): + a = layer.self_attn + for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): + lin.weight.data.normal_(0, 0.03) + for norm in (a.q_a_layernorm, a.kv_a_layernorm): + norm.weight.data.normal_(1.0, 0.02) + layer.input_layernorm.weight.data.normal_(1.0, 0.02) + layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02) + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + mlp.gate.weight.data = torch.randn( + cfg.n_routed_experts, cfg.hidden_size, device=DEVICE) + mlp.gate.e_score_correction_bias.data = torch.randn( + cfg.n_routed_experts, device=DEVICE) + mlp.experts.gate_up_proj.data.normal_(0, 0.05) + mlp.experts.down_proj.data.normal_(0, 0.05) + mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05) + else: + mlp.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.down_proj.weight.data.normal_(0, 0.05) + + +def _build_model(cfg, dtype): + model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=dtype) + model.model.embed_tokens.weight.data.normal_(0, 0.05) + model.model.norm.weight.data.normal_(1.0, 0.02) + model.lm_head.weight.data.normal_(0, 0.02) + for layer in model.model.layers: + _fill_layer(layer, cfg) + return model.eval() + + +# -------------------------------------------------------------------------- +# Test +# -------------------------------------------------------------------------- + +def test_full_forward_logits_match_reference(): + torch.manual_seed(0) + cfg = KimiK2Config.reduced() + dtype = torch.bfloat16 + model = _build_model(cfg, dtype) + # The stack spans the dense->MoE transition (first_k_dense_replace=1). + assert not isinstance(model.model.layers[0].mlp, KimiSparseMoeBlock) + assert isinstance(model.model.layers[1].mlp, KimiSparseMoeBlock) + + T = 8 + ids = torch.randint(0, cfg.vocab_size, (T,), device=DEVICE) + pos = torch.arange(T, device=DEVICE) + + cache = _MockMLACache(cfg.padded_head_dim) + with torch.no_grad(): + got = model(ids, cache, pos) + expected = _ref_forward(model, cfg, ids, pos) + + # advance_seq_lens is called exactly once per forward (after the layer loop), + # not once per layer — the cache-handle contract mirrored from Orpheus. + assert cache.advance_calls == 1 + assert got.shape == (T, cfg.vocab_size) + torch.testing.assert_close(got, expected, rtol=3e-2, atol=3e-2) diff --git a/test/integration/test_kimi_mla.py b/test/integration/test_kimi_mla.py new file mode 100644 index 000000000..602e9d8b7 --- /dev/null +++ b/test/integration/test_kimi_mla.py @@ -0,0 +1,197 @@ +"""M3 golden tests for Kimi-K2.7 MLA attention (naive/materialized path). + +Three goldens against independent references cited to vLLM: + - YARN RoPE (KimiYarnRotaryEmbedding) vs a DeepseekScalingRotaryEmbedding-style + forward_static, + - the q/k/v assembly (projections + rope-on-slice + k_pe broadcast + v-pad + + mscale^2 q-prescale) captured at ``run_attention`` via a mock cache handle, and + - the full attention forward (+ causal attention + output slice + o_proj). + +A ``_MockMLACache`` stands in for the paged cache: its ``run_attention`` does a +causal SDPA at the fixed ``1/sqrt(qk_head_dim)`` scale (what FlashInfer uses), +which is what lets us golden the MLA math without the paged engine. The real +FlashInfer path over a 192-dim cache is exercised at M4/M6. + +Refs: vLLM ``models/deepseek_v2.py::DeepseekV2Attention`` (naive path) and +``rotary_embedding/deepseek_scaling_rope.py``. + +Run: pytest test/integration/test_kimi_mla.py -v +""" +import pytest +import torch +import torch.nn.functional as F + +from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention +from mstar.model.kimi_k2_7.components.rope import ( + KimiYarnRotaryEmbedding, + _yarn_find_correction_range, + _yarn_linear_ramp_mask, + rotate_gptj, + yarn_get_mscale, +) +from mstar.model.kimi_k2_7.config import KimiK2Config + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="M3 golden tests need a GPU (MLA RMSNorm uses a FlashInfer kernel)", +) + +DEVICE = "cuda" + + +# -------------------------------------------------------------------------- +# References +# -------------------------------------------------------------------------- + +def _ref_rmsnorm(x, weight, eps): + x32 = x.float() + x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) + return weight * x32.to(x.dtype) + + +def _ref_yarn_rope(pos, q_pe, k_pe, rotary_dim, base, factor, max_pos, + beta_fast, beta_slow, mscale, mscale_all_dim): + pos_freqs = base ** (torch.arange(0, rotary_dim, 2, device=q_pe.device).float() / rotary_dim) + ext, interp = 1.0 / pos_freqs, 1.0 / (factor * pos_freqs) + low, high = _yarn_find_correction_range(beta_fast, beta_slow, rotary_dim, base, max_pos) + mask = (1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float).to(q_pe.device)) + inv_freq = interp * (1 - mask) + ext * mask + amp = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim) + freqs = torch.outer(pos.float(), inv_freq) + cos = (freqs.cos() * amp).repeat_interleave(2, -1).unsqueeze(-2) + sin = (freqs.sin() * amp).repeat_interleave(2, -1).unsqueeze(-2) + qr = q_pe.float() * cos + rotate_gptj(q_pe.float()) * sin + kr = k_pe.float() * cos + rotate_gptj(k_pe.float()) * sin + return qr.to(q_pe.dtype), kr.to(k_pe.dtype) + + +class _MockMLACache: + """Paged-cache stand-in: causal SDPA at 1/sqrt(head_dim).""" + + def __init__(self, head_dim: int): + self.scale = head_dim ** -0.5 + self.captured: dict = {} + + def set_layer_idx(self, _i): # noqa: D401 + pass + + def set_active_label(self, _l): + pass + + def advance_seq_lens(self, *_a, **_k): + pass + + def run_attention(self, q, k, v): + self.captured = {"q": q.clone(), "k": k.clone(), "v": v.clone()} + qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) # (H,T,D) + scores = torch.einsum("hqd,hkd->hqk", qt, kt) * self.scale + num_tokens = q.shape[0] + causal = torch.triu( + torch.full((num_tokens, num_tokens), float("-inf"), device=q.device), diagonal=1) + attn = (scores + causal).softmax(-1) + return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) + + +def _rope_kwargs(cfg): + r = cfg.rope_scaling + return dict( + rotary_dim=cfg.qk_rope_head_dim, base=cfg.rope_theta, factor=r["factor"], + max_pos=r["original_max_position_embeddings"], + beta_fast=r.get("beta_fast", 32), beta_slow=r.get("beta_slow", 1), + mscale=r.get("mscale", 1.0), mscale_all_dim=r.get("mscale_all_dim", 0.0), + ) + + +def _ref_mla(attn: KimiMLAAttention, cfg, h, pos, scale, boost): + """Independent MLA forward using weights extracted from ``attn``.""" + T, H = h.shape[0], attn.num_heads + Dnope, Drope, Dv, L = ( + cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank) + eps = cfg.rms_norm_eps + q = F.linear(h, attn.q_a_proj.weight) + q = _ref_rmsnorm(q, attn.q_a_layernorm.weight, eps) + q = F.linear(q, attn.q_b_proj.weight).view(T, H, cfg.qk_head_dim) + q_nope, q_pe = q.split([Dnope, Drope], dim=-1) + latent = F.linear(h, attn.kv_a_proj_with_mqa.weight) + kv_a, k_pe = latent.split([L, Drope], dim=-1) + kv_a = _ref_rmsnorm(kv_a, attn.kv_a_layernorm.weight, eps) + kv = F.linear(kv_a, attn.kv_b_proj.weight).view(T, H, Dnope + Dv) + k_nope, v = kv.split([Dnope, Dv], dim=-1) + k_pe = k_pe.view(T, 1, Drope) + rk = _rope_kwargs(cfg) + q_pe, k_pe = _ref_yarn_rope(pos, q_pe, k_pe, **rk) + # M6 mitigation: q/k assembled at Dqk then zero-padded to padded_head_dim, v + # padded from Dv to padded_head_dim (see KimiMLAAttention.forward). + pad = cfg.padded_head_dim + q = F.pad(torch.cat([q_nope, q_pe], dim=-1), [0, pad - cfg.qk_head_dim]) * boost + k = F.pad(torch.cat([k_nope, k_pe.expand(T, H, Drope)], dim=-1), [0, pad - cfg.qk_head_dim]) + v = F.pad(v, [0, pad - Dv]) + return q, k, v + + +# -------------------------------------------------------------------------- +# Tests +# -------------------------------------------------------------------------- + +def test_yarn_rope_matches_reference(): + torch.manual_seed(0) + # mscale != mscale_all_dim so the cos/sin amplitude factor is non-trivial. + rd, base, factor, max_pos = 8, 50000.0, 32.0, 4096 + ms, msad = 1.0, 0.5 + rope = KimiYarnRotaryEmbedding(rd, base, factor, max_pos, 32, 1, ms, msad).to(DEVICE) + pos = torch.arange(6, device=DEVICE) + q = torch.randn(6, 4, rd, device=DEVICE) + k = torch.randn(6, 1, rd, device=DEVICE) + + gq, gk = rope(pos, q, k) + rq, rk = _ref_yarn_rope(pos, q, k, rd, base, factor, max_pos, 32, 1, ms, msad) + torch.testing.assert_close(gq, rq, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(gk, rk, rtol=1e-4, atol=1e-4) + assert abs(rope.mscale - 1.0) > 1e-3 # amplitude path is exercised + + +def _build_attention(cfg, dtype): + attn = KimiMLAAttention(cfg).to(device=DEVICE, dtype=dtype) + for lin in (attn.q_a_proj, attn.q_b_proj, attn.kv_a_proj_with_mqa, + attn.kv_b_proj, attn.o_proj): + lin.weight.data.normal_(0, 0.03) + for norm in (attn.q_a_layernorm, attn.kv_a_layernorm): + norm.weight.data.normal_(1.0, 0.02) + return attn + + +def test_mla_qkv_assembly_matches_reference(): + torch.manual_seed(1) + cfg = KimiK2Config.reduced() + dtype = torch.bfloat16 + attn = _build_attention(cfg, dtype) + cache = _MockMLACache(cfg.padded_head_dim) + h = torch.randn(5, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.1 + pos = torch.arange(5, device=DEVICE) + + attn(h, cache, pos) # populates cache.captured with the assembled q/k/v + ref_q, ref_k, ref_v = _ref_mla(attn, cfg, h, pos, cache.scale, attn.softmax_scale_boost) + + torch.testing.assert_close(cache.captured["q"], ref_q, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(cache.captured["k"], ref_k, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(cache.captured["v"], ref_v, rtol=2e-2, atol=2e-2) + + +def test_mla_attention_forward_matches_reference(): + torch.manual_seed(2) + cfg = KimiK2Config.reduced() + dtype = torch.bfloat16 + attn = _build_attention(cfg, dtype) + cache = _MockMLACache(cfg.padded_head_dim) + h = torch.randn(7, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.1 + pos = torch.arange(7, device=DEVICE) + + got = attn(h, cache, pos) + + ref_q, ref_k, ref_v = _ref_mla(attn, cfg, h, pos, cache.scale, attn.softmax_scale_boost) + ref_attn = _MockMLACache(cfg.padded_head_dim).run_attention(ref_q, ref_k, ref_v) + ref_out = ref_attn[..., : cfg.v_head_dim].reshape(7, attn.num_heads * cfg.v_head_dim) + expected = F.linear(ref_out, attn.o_proj.weight) + + assert got.shape == (7, cfg.hidden_size) + torch.testing.assert_close(got, expected, rtol=3e-2, atol=3e-2) diff --git a/test/integration/test_kimi_mla_paged.py b/test/integration/test_kimi_mla_paged.py new file mode 100644 index 000000000..cbbbef1e9 --- /dev/null +++ b/test/integration/test_kimi_mla_paged.py @@ -0,0 +1,201 @@ +"""M6 step 1: the real paged MLA path, end-to-end, at the DeepSeek scale. + +This is the test that finally validates ``KimiMLAAttention`` over mstar's REAL +paged ``FlashInferCacheManager`` (genuine ``PagedAllocationManager`` + KV cache), +not the MockCacheHandle SDPA stand-in the M3/M4/M5 goldens use. + +It closes the M4 FlashInfer-192 blocker. The naive MLA pads q/k (from +``qk_head_dim``) and v (from ``v_head_dim``) up to ``padded_head_dim`` — the +smallest FlashInfer-SM90-supported head_dim {64,128,256} >= ``qk_head_dim`` — so +the reduced ``qk_head_dim=24`` becomes 64 (real Kimi 192 -> 256). The Hopper +prefill kernel ``static_assert``s ``head_dim_vo in {64,128,256}``, so the raw 24 +(and 192) fail to JIT-build; 64 builds and runs. + +The correctness crux is the **softmax-scale compensation**. run_attention applies +a fixed ``1/sqrt(padded_head_dim)`` scale, but DeepSeek's intended softmax scale +is ``qk_head_dim**-0.5 * mscale**2``. The zero-pad dims contribute 0 to q·k, so +we fold ``boost = mscale**2 * sqrt(padded_head_dim / qk_head_dim)`` into q: + + scores = (q*boost)·k * padded_head_dim**-0.5 + = q·k * mscale**2 * sqrt(padded/qk) * padded**-0.5 + = q·k * mscale**2 * qk**-0.5 (the DeepSeek scale). + +The reference below is the **independent DeepSeek computation** — projections + +YARN RoPE + causal SDPA at ``qk_head_dim**-0.5 * mscale**2`` over the UNPADDED q/k +(Dqk) and v (Dv), then output slice + o_proj. Matching it proves the padded paged +run + scale compensation reproduce the intended result exactly. + +Run: pytest test/integration/test_kimi_mla_paged.py -v +""" +import pytest +import torch +import torch.nn.functional as F + +from mstar.communication.tensors import LocalTransferEngine +from mstar.engine.cache_manager import WorkspaceBufferManager, create_cache_manager +from mstar.engine.kv_store import ( + KVCacheConfig, + PagedAllocationManager, + TransferEngineInfo, +) +from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention +from mstar.model.kimi_k2_7.components.rope import ( + _yarn_find_correction_range, + _yarn_linear_ramp_mask, + rotate_gptj, + yarn_get_mscale, +) +from mstar.model.kimi_k2_7.config import KimiK2Config + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="real FlashInfer paged MLA needs a GPU", +) + +DEVICE = torch.device("cuda") + + +# -------------------------------------------------------------------------- +# Real paged cache manager (mirrors test_kimi_flashinfer_attention.py). +# -------------------------------------------------------------------------- + +def _make_real_cache_manager(num_heads, head_dim, dtype, page_size=128, max_num_pages=8): + kv_cache = torch.zeros( + 2, max_num_pages, 2, page_size, num_heads, head_dim, + dtype=dtype, device=DEVICE, + ).contiguous() + kv_cfg = KVCacheConfig( + num_layers=2, num_kv_heads=num_heads, head_dim=head_dim, + max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages, + page_size=page_size, num_qo_heads=num_heads, + ) + transfer_info = TransferEngineInfo( + my_entity_id="kimi_mla_paged_test", + my_session_id="kimi_session", + transfer_engine=LocalTransferEngine("localhost"), + ) + alloc = PagedAllocationManager( + config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info, + ) + alloc.add_request("r0", ["main"]) + buffers = WorkspaceBufferManager(64 * 1024 * 1024, device=DEVICE) + cm = create_cache_manager( + request_ids=["r0"], + active_labels_per_request={"r0": "main"}, + kv_cache=kv_cache, + alloc_manager=alloc, + buffer_manager=buffers, + kv_cache_config=kv_cfg, + device=DEVICE, + ) + return cm, alloc + + +# -------------------------------------------------------------------------- +# Independent DeepSeek reference (no pad; scale = qk_head_dim**-0.5 * mscale**2). +# -------------------------------------------------------------------------- + +def _ref_rmsnorm(x, weight, eps): + x32 = x.float() + x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) + return weight * x32.to(x.dtype) + + +def _ref_yarn_rope(pos, q_pe, k_pe, rotary_dim, base, factor, max_pos, + beta_fast, beta_slow, mscale, mscale_all_dim): + pos_freqs = base ** (torch.arange(0, rotary_dim, 2, device=q_pe.device).float() / rotary_dim) + ext, interp = 1.0 / pos_freqs, 1.0 / (factor * pos_freqs) + low, high = _yarn_find_correction_range(beta_fast, beta_slow, rotary_dim, base, max_pos) + mask = 1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float).to(q_pe.device) + inv_freq = interp * (1 - mask) + ext * mask + amp = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim) + freqs = torch.outer(pos.float(), inv_freq) + cos = (freqs.cos() * amp).repeat_interleave(2, -1).unsqueeze(-2) + sin = (freqs.sin() * amp).repeat_interleave(2, -1).unsqueeze(-2) + qr = q_pe.float() * cos + rotate_gptj(q_pe.float()) * sin + kr = k_pe.float() * cos + rotate_gptj(k_pe.float()) * sin + return qr.to(q_pe.dtype), kr.to(k_pe.dtype) + + +def _sdpa_causal(q, k, v, scale): + qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) # (H,T,D) + T = q.shape[0] + causal = torch.triu( + torch.full((T, T), float("-inf"), device=q.device), diagonal=1) + attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1) + return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) + + +def _ref_deepseek_mla(attn: KimiMLAAttention, cfg, h, pos): + """The intended DeepSeek MLA output: NO padding, scale = qk**-0.5 * mscale**2.""" + T, H = h.shape[0], attn.num_heads + Dnope, Drope, Dv, L = ( + cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank) + eps = cfg.rms_norm_eps + q = _ref_rmsnorm(F.linear(h, attn.q_a_proj.weight), attn.q_a_layernorm.weight, eps) + q = F.linear(q, attn.q_b_proj.weight).view(T, H, cfg.qk_head_dim) + q_nope, q_pe = q.split([Dnope, Drope], dim=-1) + latent = F.linear(h, attn.kv_a_proj_with_mqa.weight) + kv_a, k_pe = latent.split([L, Drope], dim=-1) + kv = F.linear(_ref_rmsnorm(kv_a, attn.kv_a_layernorm.weight, eps), + attn.kv_b_proj.weight).view(T, H, Dnope + Dv) + k_nope, v = kv.split([Dnope, Dv], dim=-1) + k_pe = k_pe.view(T, 1, Drope) + r = cfg.rope_scaling + q_pe, k_pe = _ref_yarn_rope( + pos, q_pe, k_pe, Drope, cfg.rope_theta, r["factor"], + r["original_max_position_embeddings"], r.get("beta_fast", 32), + r.get("beta_slow", 1), r.get("mscale", 1.0), r.get("mscale_all_dim", 0.0)) + q = torch.cat([q_nope, q_pe], dim=-1) # (T, H, Dqk) — NOT padded + k = torch.cat([k_nope, k_pe.expand(T, H, Drope)], dim=-1) # (T, H, Dqk) + mscale = yarn_get_mscale(r["factor"], r.get("mscale_all_dim", 0.0)) + deepseek_scale = cfg.qk_head_dim ** -0.5 * mscale * mscale + out = _sdpa_causal(q, k, v, deepseek_scale) # v is Dv-wide, output Dv-wide + out = out.reshape(T, H * Dv) + return F.linear(out, attn.o_proj.weight) + + +def _build_attention(cfg, dtype): + attn = KimiMLAAttention(cfg).to(device=DEVICE, dtype=dtype) + for lin in (attn.q_a_proj, attn.q_b_proj, attn.kv_a_proj_with_mqa, + attn.kv_b_proj, attn.o_proj): + lin.weight.data.normal_(0, 0.03) + for norm in (attn.q_a_layernorm, attn.kv_a_layernorm): + norm.weight.data.normal_(1.0, 0.02) + return attn + + +def test_paged_mla_matches_deepseek_sdpa(): + """KimiMLAAttention through the REAL paged FlashInferCacheManager (head_dim = + padded_head_dim = 64) == the independent DeepSeek MLA at qk**-0.5 * mscale**2. + + This validates both (a) the real paged path builds+runs at the padded head_dim + (the M4 FlashInfer-192 blocker mitigation), and (b) the scale compensation is + exactly right — the padded run reproduces the unpadded DeepSeek scale. + """ + torch.manual_seed(0) + cfg = KimiK2Config.reduced() + assert cfg.qk_head_dim == 24 and cfg.padded_head_dim == 64 # the mitigation + dtype = torch.bfloat16 + attn = _build_attention(cfg, dtype) + + T = 6 + h = torch.randn(T, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.1 + pos = torch.arange(T, device=DEVICE) + + cm, alloc = _make_real_cache_manager(cfg.num_attention_heads, cfg.padded_head_dim, dtype) + try: + cm.set_active_label("main") + cm.plan_attention(seq_lens=[T], is_causal=True, dtype=dtype) + cm.set_layer_idx(0) + with torch.no_grad(): + got = attn(h, cm, pos) + torch.cuda.synchronize() + finally: + alloc.cleanup() + + expected = _ref_deepseek_mla(attn, cfg, h, pos) + assert got.shape == (T, cfg.hidden_size) + # bf16 through the real FlashInfer kernel; the scale compensation is exact in + # exact arithmetic, so any residual is pure bf16 rounding. + torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2) diff --git a/test/integration/test_kimi_moe.py b/test/integration/test_kimi_moe.py new file mode 100644 index 000000000..5714ff00b --- /dev/null +++ b/test/integration/test_kimi_moe.py @@ -0,0 +1,213 @@ +"""M2 golden tests for Kimi-K2.7 fine-grained MoE. + +Verifies the new DeepSeek-V3 MoE math against independent references: + - the group-limited sigmoid ``noaux_tc`` router (KimiMoEGate), + - the fused expert dispatch (reused fused-expert GEMM), and + - the full MoE block (routed + ungated shared expert). + +References are inlined (self-contained; no dependency on the local golden +harness) and cited to vLLM ``fused_moe/cpu_fused_moe.py::grouped_topk`` and +``models/deepseek_v2.py::DeepseekV2MoE``. + +GPU test: the fused expert GEMM (``fused_experts``) is CUDA/bf16-only, so the +block/dispatch tests run on ``cuda``; the suite skips without a GPU. + +Run: pytest test/integration/test_kimi_moe.py -v +""" +import pytest +import torch +import torch.nn.functional as F + +from mstar.model.components.moe import _dispatch +from mstar.model.kimi_k2_7.components.language_model import build_moe_block +from mstar.model.kimi_k2_7.components.moe import KimiMoEGate +from mstar.model.kimi_k2_7.config import KimiK2Config + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="M2 golden tests need a GPU (fused expert GEMM is CUDA/bf16-only)", +) + +DEVICE = "cuda" + + +# -------------------------------------------------------------------------- +# Independent references (cited to vLLM cpu_fused_moe.py / deepseek_v2.py) +# -------------------------------------------------------------------------- + +def _ref_grouped_topk( + logits: torch.Tensor, bias: torch.Tensor, n_group, topk_group, top_k, + norm_topk_prob, routed_scaling_factor, +): + """vLLM ``grouped_topk`` (sigmoid + noaux_tc) -> (weights, ids).""" + scores = logits.float().sigmoid() + T = scores.shape[0] + original = scores + scores = scores + bias.unsqueeze(0) + group_scores = scores.view(T, n_group, -1).topk(2, dim=-1)[0].sum(dim=-1) + group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = ( + group_mask.unsqueeze(-1) + .expand(T, n_group, scores.shape[-1] // n_group) + .reshape(T, -1) + ) + masked = scores.masked_fill(~score_mask.bool(), float("-inf")) + ids = torch.topk(masked, k=top_k, dim=-1, sorted=False)[1] + weights = original.gather(1, ids) + if norm_topk_prob: + weights = weights / weights.sum(dim=-1, keepdim=True) + weights = weights * routed_scaling_factor + return weights, ids + + +def _ref_routed_experts(h, gate_up, down, weights, ids): + """Naive per-token top-k expert loop matching ``fused_experts`` semantics.""" + T, H = h.shape + inter = down.shape[-1] + out = torch.zeros(T, H, dtype=h.dtype, device=h.device) + for t in range(T): + for j in range(ids.shape[1]): + e = int(ids[t, j]) + gu = gate_up[e] @ h[t] # (2*inter,) + g, u = gu[:inter], gu[inter:] + expert_out = down[e] @ (F.silu(g) * u) # (H,) + out[t] += weights[t, j] * expert_out + return out + + +def _ref_swiglu(x, gate_w, up_w, down_w): + return F.linear(F.silu(F.linear(x, gate_w)) * F.linear(x, up_w), down_w) + + +def _dense_combine(ids, weights, num_experts): + """Scatter (ids, weights) into a dense (T, E) vector for order-insensitive + comparison (topk with sorted=False returns experts in arbitrary order).""" + dense = torch.zeros(ids.shape[0], num_experts, device=ids.device) + dense.scatter_(1, ids, weights.float()) + return dense + + +# -------------------------------------------------------------------------- +# Router +# -------------------------------------------------------------------------- + +def test_moe_gate_matches_reference(): + torch.manual_seed(0) + cfg = KimiK2Config.reduced() + gate = KimiMoEGate( + cfg.hidden_size, cfg.n_routed_experts, cfg.num_experts_per_tok, + cfg.n_group, cfg.topk_group, cfg.routed_scaling_factor, + cfg.scoring_func, cfg.topk_method, cfg.norm_topk_prob, + ).to(DEVICE) + W = torch.randn(cfg.n_routed_experts, cfg.hidden_size, device=DEVICE) + bias = torch.randn(cfg.n_routed_experts, device=DEVICE) + gate.weight.data.copy_(W) + gate.e_score_correction_bias.data.copy_(bias) + + h = torch.randn(6, cfg.hidden_size, device=DEVICE) + got_w, got_ids = gate(h) + + logits = F.linear(h.float(), W.float()) + exp_w, exp_ids = _ref_grouped_topk( + logits, bias, cfg.n_group, cfg.topk_group, cfg.num_experts_per_tok, + cfg.norm_topk_prob, cfg.routed_scaling_factor, + ) + torch.testing.assert_close( + _dense_combine(got_ids, got_w, cfg.n_routed_experts), + _dense_combine(exp_ids, exp_w, cfg.n_routed_experts), + rtol=1e-5, atol=1e-5, + ) + + +def test_moe_gate_group_limited_routing(): + """With n_group=2/topk_group=1, every selected expert must come from the + single kept group — the crux of group-limited routing.""" + torch.manual_seed(1) + n_experts, n_group, topk_group, top_k = 8, 2, 1, 2 + experts_per_group = n_experts // n_group + gate = KimiMoEGate( + hidden_size=16, n_routed_experts=n_experts, num_experts_per_tok=top_k, + n_group=n_group, topk_group=topk_group, routed_scaling_factor=1.0, + ).to(DEVICE) + gate.weight.data.copy_(torch.randn(n_experts, 16, device=DEVICE)) + gate.e_score_correction_bias.data.copy_(torch.randn(n_experts, device=DEVICE)) + + h = torch.randn(20, 16, device=DEVICE) + _, ids = gate(h) + + # Each token's chosen experts share one group index. + groups = ids // experts_per_group + assert (groups == groups[:, :1]).all(), "experts crossed group boundary" + + +# -------------------------------------------------------------------------- +# Fused expert dispatch (trivial fixed router) +# -------------------------------------------------------------------------- + +def test_expert_dispatch_matches_naive(): + torch.manual_seed(2) + cfg = KimiK2Config.reduced() + dtype = torch.bfloat16 + T, H, I, E = 5, cfg.hidden_size, cfg.moe_intermediate_size, cfg.n_routed_experts + + h = torch.randn(T, H, device=DEVICE, dtype=dtype) * 0.1 + gate_up = torch.randn(E, 2 * I, H, device=DEVICE, dtype=dtype) * 0.05 + down = torch.randn(E, H, I, device=DEVICE, dtype=dtype) * 0.05 + # Trivial fixed router: every token -> experts {0, 1}, fixed weights. + ids = torch.tensor([[0, 1]] * T, device=DEVICE) + weights = torch.full((T, 2), 0.5, device=DEVICE, dtype=dtype) + + got = _dispatch(h, gate_up, down, E, ids, weights) + expected = _ref_routed_experts(h, gate_up, down, weights, ids) + torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2) + + +# -------------------------------------------------------------------------- +# Full MoE block (routed + ungated shared) +# -------------------------------------------------------------------------- + +def test_moe_block_matches_reference(): + torch.manual_seed(3) + cfg = KimiK2Config.reduced() + dtype = torch.bfloat16 + T, H, I, E = 7, cfg.hidden_size, cfg.moe_intermediate_size, cfg.n_routed_experts + shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts + + block = build_moe_block(cfg).to(device=DEVICE, dtype=dtype) + + gate_w = torch.randn(E, H, device=DEVICE) # fp32 router + bias = torch.randn(E, device=DEVICE) + expert_gate_up = torch.randn(E, 2 * I, H, device=DEVICE, dtype=dtype) * 0.05 + expert_down = torch.randn(E, H, I, device=DEVICE, dtype=dtype) * 0.05 + sh_gate = torch.randn(shared_inter, H, device=DEVICE, dtype=dtype) * 0.05 + sh_up = torch.randn(shared_inter, H, device=DEVICE, dtype=dtype) * 0.05 + sh_down = torch.randn(H, shared_inter, device=DEVICE, dtype=dtype) * 0.05 + + # Keep the router in fp32 (deterministic selection); load fused expert + + # shared weights. + block.gate.weight.data = gate_w + block.gate.e_score_correction_bias.data = bias + block.experts.gate_up_proj.data.copy_(expert_gate_up) + block.experts.down_proj.data.copy_(expert_down) + block.shared_expert.gate_up_proj.weight_loader( + block.shared_expert.gate_up_proj.weight, sh_gate, loaded_shard_id=0) + block.shared_expert.gate_up_proj.weight_loader( + block.shared_expert.gate_up_proj.weight, sh_up, loaded_shard_id=1) + block.shared_expert.down_proj.weight_loader( + block.shared_expert.down_proj.weight, sh_down) + + h = torch.randn(T, H, device=DEVICE, dtype=dtype) * 0.1 + got = block(h) + + logits = F.linear(h.float(), gate_w.float()) + weights, ids = _ref_grouped_topk( + logits, bias, cfg.n_group, cfg.topk_group, cfg.num_experts_per_tok, + cfg.norm_topk_prob, cfg.routed_scaling_factor, + ) + routed = _ref_routed_experts( + h, expert_gate_up, expert_down, weights.to(dtype), ids) + shared = _ref_swiglu(h, sh_gate, sh_up, sh_down) + expected = routed + shared + torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2) diff --git a/test/integration/test_kimi_submodule.py b/test/integration/test_kimi_submodule.py new file mode 100644 index 000000000..19f2cf999 --- /dev/null +++ b/test/integration/test_kimi_submodule.py @@ -0,0 +1,308 @@ +"""M6 step 6: submodule-level end-to-end through the REAL paged cache. + +This is the M6 correctness gate for serving: it exercises the whole +``KimiK2Model.get_submodule`` -> ``KimiLLMSubmodule`` -> real +``FlashInferCacheManager`` path on the reduced config with synthetic weights. + +Two things are validated: + +1. **The real build path.** ``get_submodule`` constructs ``KimiForCausalLM`` on + the meta device, casts to bf16 on meta, ``to_empty(cuda)``, and runs the M5 HF + loader — the production ``meta -> to_empty -> load_weights`` path (where the M5 + rope-buffer bug would have bitten). We assert the loaded model carries ZERO + buffers (M6 buffer audit) so no derived tensor survives as garbage. + +2. **Serving lifecycle over the paged MLA.** We drive + ``prepare_inputs -> preprocess -> forward`` through a genuine + ``FlashInferCacheManager`` (head_dim = padded_head_dim = 64) for a prefill plus + several decode steps, asserting sane token generation. The prefill logits are + checked against a mock-cache forward of the SAME loaded model at the + DeepSeek-correct scale, tying the paged serving path to the validated + MockCacheHandle goldens. + +``mstar-serve`` full-stack e2e (conductor + worker processes + SHM ports + CUDA +graph capture) is NOT run here — that infra isn't stood up in this environment. +This submodule-level test is the required correctness gate; see the M6 notes in +kimi-port-plan for the serve status. + +Run: pytest test/integration/test_kimi_submodule.py -v +""" +import pytest +import torch + +from mstar.communication.tensors import LocalTransferEngine +from mstar.engine.cache_manager import WorkspaceBufferManager, create_cache_manager +from mstar.engine.kv_store import ( + KVCacheConfig, + PagedAllocationManager, + TransferEngineInfo, +) +from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM +from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock +from mstar.model.kimi_k2_7.config import KimiK2Config +from mstar.model.kimi_k2_7.kimi_model import KimiK2Model +from mstar.model.kimi_k2_7.submodules import KimiLLMSubmodule +from mstar.model.submodule_base import ModelInputsFromEngine + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="submodule e2e needs a GPU (real FlashInfer paged cache)", +) + +DEVICE = torch.device("cuda") + + +# -------------------------------------------------------------------------- +# Synthetic HF DeepSeek-V3 checkpoint (same serialization as +# test_kimi_weight_loading — un-fuse every fused param back to HF keys). +# -------------------------------------------------------------------------- + +def _fill_layer(layer, cfg): + a = layer.self_attn + for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): + lin.weight.data.normal_(0, 0.03) + for norm in (a.q_a_layernorm, a.kv_a_layernorm): + norm.weight.data.normal_(1.0, 0.02) + layer.input_layernorm.weight.data.normal_(1.0, 0.02) + layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02) + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + mlp.gate.weight.data.normal_(0, 1) + mlp.gate.e_score_correction_bias.data = torch.randn( + cfg.n_routed_experts, device=DEVICE, dtype=torch.float32) + mlp.experts.gate_up_proj.data.normal_(0, 0.05) + mlp.experts.down_proj.data.normal_(0, 0.05) + mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05) + else: + mlp.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.down_proj.weight.data.normal_(0, 0.05) + + +def _build_reference(cfg): + model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=torch.bfloat16) + model.model.embed_tokens.weight.data.normal_(0, 0.05) + model.model.norm.weight.data.normal_(1.0, 0.02) + model.lm_head.weight.data.normal_(0, 0.02) + for layer in model.model.layers: + _fill_layer(layer, cfg) + return model.eval() + + +def _hf_checkpoint(model, cfg): + inter = cfg.intermediate_size + moe_inter = cfg.moe_intermediate_size + shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts + m = model.model + sd = {"model.embed_tokens.weight": m.embed_tokens.weight} + for i, layer in enumerate(m.layers): + p = f"model.layers.{i}." + a = layer.self_attn + sd[p + "self_attn.q_a_proj.weight"] = a.q_a_proj.weight + sd[p + "self_attn.q_a_layernorm.weight"] = a.q_a_layernorm.weight + sd[p + "self_attn.q_b_proj.weight"] = a.q_b_proj.weight + sd[p + "self_attn.kv_a_proj_with_mqa.weight"] = a.kv_a_proj_with_mqa.weight + sd[p + "self_attn.kv_a_layernorm.weight"] = a.kv_a_layernorm.weight + sd[p + "self_attn.kv_b_proj.weight"] = a.kv_b_proj.weight + sd[p + "self_attn.o_proj.weight"] = a.o_proj.weight + sd[p + "input_layernorm.weight"] = layer.input_layernorm.weight + sd[p + "post_attention_layernorm.weight"] = layer.post_attention_layernorm.weight + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + sd[p + "mlp.gate.weight"] = mlp.gate.weight + sd[p + "mlp.gate.e_score_correction_bias"] = mlp.gate.e_score_correction_bias + gup, dwn = mlp.experts.gate_up_proj, mlp.experts.down_proj + for e in range(cfg.n_routed_experts): + sd[p + f"mlp.experts.{e}.gate_proj.weight"] = gup[e, :moe_inter, :] + sd[p + f"mlp.experts.{e}.up_proj.weight"] = gup[e, moe_inter:, :] + sd[p + f"mlp.experts.{e}.down_proj.weight"] = dwn[e] + sh = mlp.shared_expert + sd[p + "mlp.shared_experts.gate_proj.weight"] = sh.gate_up_proj.weight[:shared_inter] + sd[p + "mlp.shared_experts.up_proj.weight"] = sh.gate_up_proj.weight[shared_inter:] + sd[p + "mlp.shared_experts.down_proj.weight"] = sh.down_proj.weight + else: + sd[p + "mlp.gate_proj.weight"] = mlp.gate_up_proj.weight[:inter] + sd[p + "mlp.up_proj.weight"] = mlp.gate_up_proj.weight[inter:] + sd[p + "mlp.down_proj.weight"] = mlp.down_proj.weight + sd["model.norm.weight"] = m.norm.weight + sd["lm_head.weight"] = model.lm_head.weight + return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()} + + +# -------------------------------------------------------------------------- +# Real paged cache + mock cache (DeepSeek-correct scale via padded_head_dim). +# -------------------------------------------------------------------------- + +def _make_real_cache_manager(cfg, dtype, page_size=128, max_num_pages=8): + num_heads = cfg.num_attention_heads + head_dim = cfg.padded_head_dim + kv_cache = torch.zeros( + cfg.num_hidden_layers, max_num_pages, 2, page_size, num_heads, head_dim, + dtype=dtype, device=DEVICE, + ).contiguous() + kv_cfg = KVCacheConfig( + num_layers=cfg.num_hidden_layers, num_kv_heads=num_heads, head_dim=head_dim, + max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages, + page_size=page_size, num_qo_heads=num_heads, + ) + transfer_info = TransferEngineInfo( + my_entity_id="kimi_submodule_test", my_session_id="kimi_session", + transfer_engine=LocalTransferEngine("localhost"), + ) + alloc = PagedAllocationManager( + config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info) + alloc.add_request("r0", ["main"]) + buffers = WorkspaceBufferManager(64 * 1024 * 1024, device=DEVICE) + cm = create_cache_manager( + request_ids=["r0"], active_labels_per_request={"r0": "main"}, + kv_cache=kv_cache, alloc_manager=alloc, buffer_manager=buffers, + kv_cache_config=kv_cfg, device=DEVICE, + ) + return cm, alloc + + +def _sdpa_causal(q, k, v, scale): + qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) + T = q.shape[0] + causal = torch.triu(torch.full((T, T), float("-inf"), device=q.device), diagonal=1) + attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1) + return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) + + +class _MockMLACache: + def __init__(self, head_dim): + self.scale = head_dim ** -0.5 + + def set_layer_idx(self, _i): + pass + + def advance_seq_lens(self, *_a, **_k): + pass + + def run_attention(self, q, k, v): + return _sdpa_causal(q, k, v, self.scale) + + +def _make_model(cfg, checkpoint_dir) -> KimiK2Model: + """A KimiK2Model wired to the synthetic checkpoint without a tokenizer. + + object.__new__ skips __init__ (which would pull a tokenizer / the full config), + so we set only what get_submodule needs — mirroring the modular test builder. + """ + model = object.__new__(KimiK2Model) + model.config = cfg + model.model_path_hf = str(checkpoint_dir) + model.cache_dir = None + model._submodule_cache = {} + return model + + +# -------------------------------------------------------------------------- +# Minimal engine-inputs + lifecycle driver. +# -------------------------------------------------------------------------- + +def _engine_inputs(cm): + return ModelInputsFromEngine( + request_ids=["r0"], per_request_info={}, cache_manager=cm, + ) + + +def _step(submodule, cm, graph_walk, token_ids): + """Drive prepare_inputs -> preprocess -> forward for one packed request.""" + engine_inputs = _engine_inputs(cm) + ar_in = submodule.prepare_inputs( + graph_walk=graph_walk, fwd_info=None, + inputs={"text_inputs": [token_ids]}, + ) + packed = submodule.preprocess(graph_walk, engine_inputs, [ar_in]) + with torch.no_grad(): + out = submodule.forward(graph_walk, engine_inputs, **packed) + return out["logits"][0], packed # (1, vocab), packed dict + + +# -------------------------------------------------------------------------- +# Test +# -------------------------------------------------------------------------- + +def test_submodule_prefill_decode_over_real_paged_cache(tmp_path): + from safetensors.torch import save_file + + torch.manual_seed(0) + cfg = KimiK2Config.reduced() + ref = _build_reference(cfg) + save_file(_hf_checkpoint(ref, cfg), str(tmp_path / "model.safetensors")) + + # --- the real build path: meta -> to(bf16) -> to_empty(cuda) -> load --- + model = _make_model(cfg, tmp_path) + submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16) + assert isinstance(submodule, KimiLLMSubmodule) + # get_submodule caches the built submodule. + assert model.get_submodule("LLM") is submodule + # M6 buffer audit: no derived tensor buffer survived the load path as garbage. + assert list(submodule.language_model.named_buffers()) == [] + p = next(submodule.language_model.parameters()) + assert p.device.type == "cuda" and p.dtype == torch.bfloat16 + + # --- prefill over the real paged FlashInfer cache (head_dim = 64) --- + T = 6 + prompt = torch.randint(0, cfg.vocab_size, (T,), device=DEVICE) + cm, alloc = _make_real_cache_manager(cfg, torch.bfloat16) + try: + prefill_logits, _ = _step(submodule, cm, "prefill", prompt) + assert prefill_logits.shape == (1, cfg.vocab_size) + assert torch.isfinite(prefill_logits).all() + + # Reference: the SAME loaded model through the mock cache at the + # DeepSeek-correct scale (padded_head_dim). Ties the paged serving path to + # the validated MockCacheHandle goldens. Loose bf16 tolerance (2-layer stack + # through the real FlashInfer kernel). + pos = torch.arange(T, device=DEVICE) + with torch.no_grad(): + ref_hidden = submodule.language_model.model( + prompt, _MockMLACache(cfg.padded_head_dim), pos) + ref_logits = submodule.lm_head(ref_hidden[-1:]) + torch.testing.assert_close(prefill_logits, ref_logits, rtol=5e-2, atol=5e-2) + + # --- a few decode steps over the accumulating paged KV cache --- + next_token = prefill_logits.argmax(-1) # (1,) + generated = [int(next_token.item())] + assert 0 <= generated[-1] < cfg.vocab_size + for _ in range(4): + logits, _ = _step(submodule, cm, "decode", next_token) + assert logits.shape == (1, cfg.vocab_size) + assert torch.isfinite(logits).all() + next_token = logits.argmax(-1) + tok = int(next_token.item()) + assert 0 <= tok < cfg.vocab_size + generated.append(tok) + finally: + alloc.cleanup() + + # Sane generation: right length, all valid ids. + assert len(generated) == 5 + assert all(0 <= t < cfg.vocab_size for t in generated) + + +def test_submodule_paged_decode_is_deterministic(tmp_path): + """Same prompt + fresh cache -> identical first token (paged path is stable and + the load is reproducible). Cheap guard against nondeterministic KV writes.""" + from safetensors.torch import save_file + + torch.manual_seed(1) + cfg = KimiK2Config.reduced() + ref = _build_reference(cfg) + save_file(_hf_checkpoint(ref, cfg), str(tmp_path / "model.safetensors")) + model = _make_model(cfg, tmp_path) + submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16) + + T = 5 + prompt = torch.randint(0, cfg.vocab_size, (T,), device=DEVICE) + tokens = [] + for _ in range(2): + cm, alloc = _make_real_cache_manager(cfg, torch.bfloat16) + try: + logits, _ = _step(submodule, cm, "prefill", prompt) + tokens.append(int(logits.argmax(-1).item())) + finally: + alloc.cleanup() + assert tokens[0] == tokens[1] diff --git a/test/integration/test_kimi_weight_loading.py b/test/integration/test_kimi_weight_loading.py new file mode 100644 index 000000000..08b2849a7 --- /dev/null +++ b/test/integration/test_kimi_weight_loading.py @@ -0,0 +1,235 @@ +"""M5 weight-loading golden test for Kimi-K2.7 / DeepSeek-V3 (synthetic checkpoint). + +The real 1T checkpoint is absent and would not fit, so this validates the loader +on a SYNTHETIC ``KimiK2Config.reduced()`` model with random bf16 weights: + + 1. build a ``KimiForCausalLM`` reference, fill it with random weights (router + ``e_score_correction_bias`` kept fp32); + 2. serialize it to a temp dir as an **HF DeepSeek-V3 checkpoint** — the exact + inverse of the loader's remap: per-expert ``gate_up_proj`` un-fused back to + ``experts.{e}.{gate,up}_proj``, dense/shared merged gate/up un-fused, singular + ``shared_expert`` -> HF plural ``shared_experts`` — as ``model.safetensors``; + 3. build a *fresh* model on ``meta`` -> ``to(bf16)`` -> ``to_empty(cuda)`` (the + production path), then load the checkpoint via the standard + ``mstar.model.loader.load_weights(model, dir, device)`` driver, which invokes + ``KimiForCausalLM.load_weights`` -> the M5 stacked rules + name remap; + 4. assert (a) every param of the loaded model equals the reference source + (exact), with targeted fused-param slice checks proving the gate/up/down + fusion for BOTH a dense (layer 0) and a MoE (layer 1) layer, and + (b) a full forward on the loaded model matches a forward on the reference + model (same mock-cache path as ``test_kimi_forward.py``). + +Confirms MLA loads strictly by name (no q_a/kv_a fusion): the attention params +appear identically in checkpoint and module, and the round-trip is exact. + +Run: pytest test/integration/test_kimi_weight_loading.py -v +""" +import pytest +import torch + +from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM +from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock +from mstar.model.kimi_k2_7.config import KimiK2Config +from mstar.model.loader import load_weights as driver_load_weights + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="M5 weight-loading golden needs a GPU (RMSNorm + fused expert GEMM)", +) + +DEVICE = "cuda" + + +# -------------------------------------------------------------------------- +# Mock paged cache (causal SDPA at 1/sqrt(head_dim)) — same as test_kimi_forward. +# -------------------------------------------------------------------------- + +def _sdpa_causal(q, k, v, scale): + qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) + T = q.shape[0] + causal = torch.triu( + torch.full((T, T), float("-inf"), device=q.device), diagonal=1) + attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1) + return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) + + +class _MockMLACache: + def __init__(self, head_dim): + self.scale = head_dim ** -0.5 + + def set_layer_idx(self, _i): + pass + + def advance_seq_lens(self, *_a, **_k): + pass + + def run_attention(self, q, k, v): + return _sdpa_causal(q, k, v, self.scale) + + +# -------------------------------------------------------------------------- +# Random weight init (router bias kept fp32) + HF-checkpoint serialization. +# -------------------------------------------------------------------------- + +def _fill_layer(layer, cfg): + a = layer.self_attn + for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): + lin.weight.data.normal_(0, 0.03) + for norm in (a.q_a_layernorm, a.kv_a_layernorm): + norm.weight.data.normal_(1.0, 0.02) + layer.input_layernorm.weight.data.normal_(1.0, 0.02) + layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02) + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + # In place so the router weight keeps the model dtype (bf16). + mlp.gate.weight.data.normal_(0, 1) + # Router selection bias stays fp32 even in a bf16 model. + mlp.gate.e_score_correction_bias.data = torch.randn( + cfg.n_routed_experts, device=DEVICE, dtype=torch.float32) + mlp.experts.gate_up_proj.data.normal_(0, 0.05) + mlp.experts.down_proj.data.normal_(0, 0.05) + mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05) + else: + mlp.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.down_proj.weight.data.normal_(0, 0.05) + + +def _build_reference(cfg): + model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=torch.bfloat16) + model.model.embed_tokens.weight.data.normal_(0, 0.05) + model.model.norm.weight.data.normal_(1.0, 0.02) + model.lm_head.weight.data.normal_(0, 0.02) + for layer in model.model.layers: + _fill_layer(layer, cfg) + return model.eval() + + +def _hf_checkpoint(model, cfg): + """Serialize the reference model to HF DeepSeek-V3 keys (inverse of the loader). + + Un-fuses every fused param back to the per-projection / per-expert checkpoint + layout so the loader has real fusion work to do. + """ + inter = cfg.intermediate_size + moe_inter = cfg.moe_intermediate_size + shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts + m = model.model + sd = {"model.embed_tokens.weight": m.embed_tokens.weight} + for i, layer in enumerate(m.layers): + p = f"model.layers.{i}." + a = layer.self_attn + # MLA — identity keys, no fusion. + sd[p + "self_attn.q_a_proj.weight"] = a.q_a_proj.weight + sd[p + "self_attn.q_a_layernorm.weight"] = a.q_a_layernorm.weight + sd[p + "self_attn.q_b_proj.weight"] = a.q_b_proj.weight + sd[p + "self_attn.kv_a_proj_with_mqa.weight"] = a.kv_a_proj_with_mqa.weight + sd[p + "self_attn.kv_a_layernorm.weight"] = a.kv_a_layernorm.weight + sd[p + "self_attn.kv_b_proj.weight"] = a.kv_b_proj.weight + sd[p + "self_attn.o_proj.weight"] = a.o_proj.weight + sd[p + "input_layernorm.weight"] = layer.input_layernorm.weight + sd[p + "post_attention_layernorm.weight"] = layer.post_attention_layernorm.weight + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + sd[p + "mlp.gate.weight"] = mlp.gate.weight + sd[p + "mlp.gate.e_score_correction_bias"] = mlp.gate.e_score_correction_bias + gup, dwn = mlp.experts.gate_up_proj, mlp.experts.down_proj + for e in range(cfg.n_routed_experts): + sd[p + f"mlp.experts.{e}.gate_proj.weight"] = gup[e, :moe_inter, :] + sd[p + f"mlp.experts.{e}.up_proj.weight"] = gup[e, moe_inter:, :] + sd[p + f"mlp.experts.{e}.down_proj.weight"] = dwn[e] + sh = mlp.shared_expert + sd[p + "mlp.shared_experts.gate_proj.weight"] = sh.gate_up_proj.weight[:shared_inter] + sd[p + "mlp.shared_experts.up_proj.weight"] = sh.gate_up_proj.weight[shared_inter:] + sd[p + "mlp.shared_experts.down_proj.weight"] = sh.down_proj.weight + else: + sd[p + "mlp.gate_proj.weight"] = mlp.gate_up_proj.weight[:inter] + sd[p + "mlp.up_proj.weight"] = mlp.gate_up_proj.weight[inter:] + sd[p + "mlp.down_proj.weight"] = mlp.down_proj.weight + sd["model.norm.weight"] = m.norm.weight + sd["lm_head.weight"] = model.lm_head.weight + # Clone to cpu + break storage aliasing (safetensors rejects shared storage). + return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()} + + +def _build_loaded(cfg, checkpoint_dir): + """Production path: meta -> to(bf16) -> to_empty(cuda) -> load_weights.""" + with torch.device("meta"): + model = KimiForCausalLM(cfg) + model = model.to(torch.bfloat16) + model.to_empty(device=DEVICE) + loaded = driver_load_weights(model, checkpoint_dir, device=DEVICE) + return model.eval(), loaded + + +# -------------------------------------------------------------------------- +# Test +# -------------------------------------------------------------------------- + +def test_weight_loading_roundtrip_and_forward(tmp_path): + from safetensors.torch import save_file + + torch.manual_seed(0) + cfg = KimiK2Config.reduced() + ref = _build_reference(cfg) + # The stack spans the dense->MoE transition (first_k_dense_replace=1). + assert not isinstance(ref.model.layers[0].mlp, KimiSparseMoeBlock) + assert isinstance(ref.model.layers[1].mlp, KimiSparseMoeBlock) + + save_file(_hf_checkpoint(ref, cfg), str(tmp_path / "model.safetensors")) + model, loaded = _build_loaded(cfg, tmp_path) + + # (a0) completeness: every param received exactly one tensor. + all_params = set(dict(model.named_parameters()).keys()) + assert loaded == all_params, ( + f"unloaded: {all_params - loaded}; spurious: {loaded - all_params}") + + # (a1) every loaded param equals the reference source, bit for bit. + ref_sd = dict(ref.named_parameters()) + for name, param in model.named_parameters(): + assert torch.equal(param, ref_sd[name]), f"mismatch at {name}" + + # (a2) router bias preserved fp32 even in a bf16 model. + bias = model.model.layers[1].mlp.gate.e_score_correction_bias + assert bias.dtype == torch.float32 + + # (a2b) regression guard (M6 buffer audit): NO derived tensor buffer survives + # meta -> to_empty as uninitialized garbage. The M6 audit of every Kimi + # submodule (attention/moe/decoder_layer/causal_lm/rope/language_model) found + # exactly one derived non-parameter tensor — the rope inv_freq — and it is + # computed lazily (M5 fix) rather than as an __init__ buffer, so the loaded + # model carries ZERO buffers. Any future __init__-computed buffer that is not + # in the checkpoint would fail this and silently corrupt the forward. + buffer_names = {n for n, _ in model.named_buffers()} + assert buffer_names == set(), f"unexpected buffers survived the load path: {buffer_names}" + + # (a3) targeted fusion checks — dense layer 0 (merged gate/up) ... + inter = cfg.intermediate_size + d_gup = model.model.layers[0].mlp.gate_up_proj.weight + r_gup = ref.model.layers[0].mlp.gate_up_proj.weight + assert torch.equal(d_gup[:inter], r_gup[:inter]) # gate half + assert torch.equal(d_gup[inter:], r_gup[inter:]) # up half + # ... and MoE layer 1 (per-expert w13 gate|up + w2 down). + mi = cfg.moe_intermediate_size + l_gup = model.model.layers[1].mlp.experts.gate_up_proj + r_egup = ref.model.layers[1].mlp.experts.gate_up_proj + r_edwn = ref.model.layers[1].mlp.experts.down_proj + for e in range(cfg.n_routed_experts): + assert torch.equal(l_gup[e, :mi], r_egup[e, :mi]) # gate:e + assert torch.equal(l_gup[e, mi:], r_egup[e, mi:]) # up:e + assert torch.equal(model.model.layers[1].mlp.experts.down_proj, r_edwn) + + # (b) full forward on the loaded model matches the reference model's forward. + # With bit-identical params (a1) AND a correctly-initialized rope (a2b), and + # since these kernels are per-instance deterministic (a repeated forward is + # bit-reproducible), the two forwards are bit-identical. The tiny bound below + # is off any tolerance boundary by ~3 orders of magnitude (measured diff is + # exactly 0.0 across runs) while still catching a gross mis-load (O(0.1+)). + T = 8 + ids = torch.randint(0, cfg.vocab_size, (T,), device=DEVICE) + pos = torch.arange(T, device=DEVICE) + with torch.no_grad(): + got = model(ids, _MockMLACache(cfg.padded_head_dim), pos) + expected = ref(ids, _MockMLACache(cfg.padded_head_dim), pos) + assert got.shape == (T, cfg.vocab_size) + torch.testing.assert_close(got, expected, rtol=1e-3, atol=1e-3) diff --git a/test/modular/test_kimi_model.py b/test/modular/test_kimi_model.py new file mode 100644 index 000000000..0cca57cc7 --- /dev/null +++ b/test/modular/test_kimi_model.py @@ -0,0 +1,106 @@ +"""M0 scaffold tests for Kimi-K2.7 (text backbone). + +Dummy mode: the model is built via ``object.__new__`` (no tokenizer, no weights, +no GPU) and only the ``Model`` contract is exercised — the graph, engine types, +KV-cache dims, and the prefill→decode→done state machine. This validates the +serving plumbing in isolation before any MLA/MoE compute exists, exactly as +``docs/adding_models.rst`` prescribes for a new model. +""" +import sys + +sys.path.insert(0, ".") + +from mstar.conductor.request_info import CurrentForwardConductorMetadata +from mstar.engine.base import EngineType +from mstar.graph.base import Loop +from mstar.model.kimi_k2_7.config import KimiK2Config +from mstar.model.kimi_k2_7.kimi_model import KimiK2Model + + +def _make_model() -> KimiK2Model: + model = object.__new__(KimiK2Model) + model.config = KimiK2Config.reduced() + model._submodule_cache = {} + return model + + +def test_kimi_graph_walks_and_engine_types(): + model = _make_model() + + walks = model.get_graph_walk_graphs() + assert set(walks) == {"prefill", "decode"} + assert isinstance(walks["decode"], Loop) + assert walks["decode"].name == "decode_loop" + + assert model.get_node_engine_types() == {"LLM": EngineType.KV_CACHE} + + +def test_kimi_kv_cache_config_matches_reduced_mla_dims(): + model = _make_model() + cfg = model.config + + kv = model.get_kv_cache_config() + assert len(kv) == 1 + (kv,) = kv + + assert kv.num_layers == cfg.num_hidden_layers == 2 + # Naive/materialized MLA: KV heads == query heads. + assert kv.num_kv_heads == cfg.num_attention_heads == 4 + assert kv.num_qo_heads == cfg.num_attention_heads == 4 + # M6 FlashInfer-SM90 mitigation: q/k/v are zero-padded from qk_head_dim (24) + # up to the smallest supported head_dim {64,128,256} >= qk_head_dim, so the + # paged cache stores head_dim == padded_head_dim == 64 (not the raw 24, which + # the Hopper prefill kernel static_asserts against). + assert cfg.qk_head_dim == cfg.qk_nope_head_dim + cfg.qk_rope_head_dim == 24 + assert kv.head_dim == cfg.padded_head_dim == 64 + assert kv.max_seq_len == cfg.max_position_embeddings + + +def test_kimi_prefill_transitions_to_decode(): + model = _make_model() + metadata = CurrentForwardConductorMetadata( + input_modalities=["text"], + output_modalities=["text"], + graph_walk="prefill", + is_prefill=True, + ) + + result = model.get_partition_forward_pass_args( + partition_name="default", + partition_metadata=metadata, + persist_signals={"new_token": []}, + ) + + assert result.full_metadata.graph_walk == "decode" + assert result.full_metadata.is_prefill is False + assert result.step_metadata["is_prefill"] is False + assert result.request_done is False + + +def test_kimi_decode_completion_marks_done(): + model = _make_model() + metadata = CurrentForwardConductorMetadata( + input_modalities=["text"], + output_modalities=["text"], + graph_walk="decode", + is_prefill=False, + ) + + result = model.get_partition_forward_pass_args( + partition_name="default", + partition_metadata=metadata, + persist_signals={}, + ) + + assert result.request_done is True + assert result.full_metadata.kwargs["decode_finished"] is True + + +def test_kimi_get_submodule_is_dummy_mode(): + model = _make_model() + # M6: get_submodule is the real meta->to_empty->load_weights build, but it + # returns None in dummy mode when no checkpoint is resolvable. _make_model sets + # no model_path_hf, so _resolve_checkpoint() -> None -> dummy mode, letting the + # modular graph tests run without a GPU or weights. + assert getattr(model, "model_path_hf", None) is None + assert model.get_submodule("LLM") is None From 305dc6fb01b0c0a8b9831571074c911d203bb85d Mon Sep 17 00:00:00 2001 From: Garv Ghai Date: Sun, 19 Jul 2026 16:13:53 +0000 Subject: [PATCH 2/9] Kimi-K2.7: serve e2e + TP correctness (reduced config) Builds on the DeepSeek-V3 text backbone to make the reduced/synthetic config both servable end-to-end and correct under tensor parallelism. Golden-verified on GPU; no shared code touched. INT4/fp8, TP8 and the real 1T checkpoint remain deferred (Phase 4). Serve (single-GPU, reduced/synthetic): - kimi_model.py: model_kwargs hook (checkpoint_path / config_variant / tokenizer_mode) lets a serving YAML point the model at a local reduced checkpoint with a UTF-8 byte tokenizer (reduced vocab_size=256). Full-size default behaviour unchanged. - config.py + submodules.py: CUDA-graph prefill capture grid is config-driven (KimiK2Config.prefill_token_buckets / prefill_capture_batch_sizes; reduced() uses a tiny [64]/[1] grid). No env vars. - configs/kimi_k2_7_repro.yaml + test/integration/test_kimi_serve_e2e.py. Verified: live mstar-serve (api_server -> conductor -> worker -> KV_CACHE engine -> decode loop) streams tokens; deterministic in-process gate. Tensor parallelism (tp=2, reduced config): - components/attention.py: MLA head-sharding (rank-local head count; column-parallel q_b/kv_b + row-parallel o_proj; latent down-projs replicated). - components/moe.py: MoE intermediate-sharding (fused_experts reduce_results + all-reduce), mirroring ParallelSparseMoeBlock; router replicated. tp=1 is byte-identical. - configs/kimi_k2_7_tp2.yaml + test/integration/test_kimi_tp.py. Verified: tp=2 == tp=1 (real 2-GPU NCCL + in-process rank-simulation goldens). --- configs/kimi_k2_7_repro.yaml | 40 ++ configs/kimi_k2_7_tp2.yaml | 39 ++ mstar/model/kimi_k2_7/components/attention.py | 34 +- mstar/model/kimi_k2_7/components/moe.py | 90 ++++- mstar/model/kimi_k2_7/config.py | 14 + mstar/model/kimi_k2_7/kimi_model.py | 38 +- mstar/model/kimi_k2_7/submodules.py | 16 +- test/integration/test_kimi_serve_e2e.py | 313 ++++++++++++++++ test/integration/test_kimi_tp.py | 351 ++++++++++++++++++ 9 files changed, 911 insertions(+), 24 deletions(-) create mode 100644 configs/kimi_k2_7_repro.yaml create mode 100644 configs/kimi_k2_7_tp2.yaml create mode 100644 test/integration/test_kimi_serve_e2e.py create mode 100644 test/integration/test_kimi_tp.py diff --git a/configs/kimi_k2_7_repro.yaml b/configs/kimi_k2_7_repro.yaml new file mode 100644 index 000000000..f284251eb --- /dev/null +++ b/configs/kimi_k2_7_repro.yaml @@ -0,0 +1,40 @@ +model: "kimi_k2_7" +# Kimi-K2.7 text backbone — REDUCED / synthetic single-GPU repro config (Phase 2 +# gap-4 serve bring-up). It drives the full serving path (API server -> conductor +# -> worker -> KV_CACHE engine -> decode loop -> tokens) on a tiny model that runs +# without the 1T checkpoint. NOT a real deployment — for that use kimi_k2_7.yaml. +# +# `model_kwargs` (forwarded to KimiK2Model.__init__) redirect this model at a +# local synthetic checkpoint + reduced config + a trivial byte tokenizer, so no +# 1T weights and no real Kimi tokenizer are needed: +# * config_variant: reduced -> KimiK2Config.reduced() (vocab 256, 2 layers) +# * checkpoint_path -> the dir written by +# tools/kimi_goldens/make_repro_checkpoint.py +# * tokenizer_mode: byte -> UTF-8 byte identity tokenizer (ids in [0,256)) +# +# Generate the checkpoint first (writes tools/kimi_goldens/repro/checkpoint, which +# is gitignored): +# python tools/kimi_goldens/make_repro_checkpoint.py +# `checkpoint_path` is relative to the mstar package root — run `mstar-serve` from +# there (the launch harness does). Make it absolute if you launch from elsewhere. +max_seq_len: 512 +model_kwargs: + config_variant: reduced + checkpoint_path: tools/kimi_goldens/repro/checkpoint + tokenizer_mode: byte +# Small paged KV cache — the reduced model needs almost nothing (2 layers, 4 heads, +# head_dim 64). 256 pages * 128 = 32768-token capacity, ~33 MB. Sized to cover +# CUDA-graph decode capture (default batch sizes up to 64, double-buffered), not +# just the single serving request. +kv_cache: + max_num_pages: 256 + page_size: 128 +# The CUDA-graph prefill capture grid is trimmed to a single short-prompt bucket +# at batch size 1 by KimiK2Config.reduced() (its prefill_token_buckets / +# prefill_capture_batch_sizes), which `config_variant: reduced` above selects — so +# no env vars are needed. The full 6x5 grid is only captured for the full model. +node_groups: + - node_names: [LLM] + ranks: [0] + tp_size: 1 + graph_walks: [prefill, decode] diff --git a/configs/kimi_k2_7_tp2.yaml b/configs/kimi_k2_7_tp2.yaml new file mode 100644 index 000000000..04de97a0b --- /dev/null +++ b/configs/kimi_k2_7_tp2.yaml @@ -0,0 +1,39 @@ +model: "kimi_k2_7" +# Kimi-K2.7 text backbone — REDUCED / synthetic TP=2 repro config (Phase 3 gap-3 +# TP correctness). Same tiny synthetic model as kimi_k2_7_repro.yaml, but the LLM +# node runs tensor-parallel across 2 ranks so the MLA head-shard + MoE +# intermediate-shard paths execute end-to-end. NOT a real deployment — the 1T +# model needs TP8 / multi-node (see kimi_k2_7.yaml). +# +# `model_kwargs` (forwarded to KimiK2Model.__init__) redirect this model at a +# local synthetic checkpoint + reduced config + a trivial byte tokenizer: +# * config_variant: reduced -> KimiK2Config.reduced() (vocab 256, 2 layers, +# 4 attention heads -> 2/rank, moe_inter 64 -> 32/rank) +# * checkpoint_path -> the dir written by +# tools/kimi_goldens/make_repro_checkpoint.py +# * tokenizer_mode: byte -> UTF-8 byte identity tokenizer (ids in [0,256)) +# +# Generate the checkpoint first (writes tools/kimi_goldens/repro/checkpoint, which +# is gitignored): +# python tools/kimi_goldens/make_repro_checkpoint.py +# `checkpoint_path` is relative to the mstar package root — run `mstar-serve` from +# there. Make it absolute if you launch from elsewhere. +# +# On this cluster (coriander) RDMA is unavailable, so launch with a non-RDMA +# tensor transport, e.g. TENSOR_PROTOCOL=SHM (single node) — see mstar/.sample.env. +max_seq_len: 512 +model_kwargs: + config_variant: reduced + checkpoint_path: tools/kimi_goldens/repro/checkpoint + tokenizer_mode: byte +kv_cache: + max_num_pages: 256 + page_size: 128 +# The LLM node is the tensor-parallel node (mirrors configs/orpheus_tp2.yaml): +# ranks [0, 1] with tp_size 2 for both the prefill and decode walks. The model +# code is TP-agnostic — the per-rank shard degree comes from here. +node_groups: + - node_names: [LLM] + ranks: [0, 1] + tp_size: 2 + graph_walks: [prefill, decode] diff --git a/mstar/model/kimi_k2_7/components/attention.py b/mstar/model/kimi_k2_7/components/attention.py index 8141f9eb5..30a63d3cf 100644 --- a/mstar/model/kimi_k2_7/components/attention.py +++ b/mstar/model/kimi_k2_7/components/attention.py @@ -23,8 +23,11 @@ Cache config for this node: ``num_kv_heads == num_qo_heads == num_attention_heads``, ``head_dim == padded_head_dim`` (256 for the real Dqk=192, 64 for the reduced -Dqk=24). Weight-absorbed MLA (native latent dims, no pad) and the -``fused_qkv_a_proj`` weight fusion are deferred to the perf backlog. +Dqk=24). Under tensor parallelism each rank materializes only its +``num_attention_heads // tp_size`` local heads (K/V and Q both shard on the head +axis — there is no separate KV-head group in the naive path), and the paged cache +reports the matching per-rank count. Weight-absorbed MLA (native latent dims, no +pad) and the ``fused_qkv_a_proj`` weight fusion are deferred to the perf backlog. """ from __future__ import annotations @@ -50,8 +53,27 @@ def __init__(self, config: KimiK2Config, comm_group: CommGroup | None = None) -> if comm_group is None: comm_group = CommGroup.trivial() - # TODO(M6): shard num_heads across TP; naive path assumes local == total. - self.num_heads = config.num_attention_heads + # MLA shards on the head dim under TP, mirroring vLLM + # ``DeepseekV2MLAAttention``: the query/kv UP-projections (``q_b_proj`` / + # ``kv_b_proj``) are ColumnParallel and ``o_proj`` is RowParallel, so each + # rank owns a contiguous block of ``num_heads // tp_size`` attention heads. + # The latent DOWN-projections (``q_a_proj`` / ``kv_a_proj_with_mqa``) and + # their RMSNorms are REPLICATED (small shared latent, no head structure). + # ``num_heads`` below is this rank's LOCAL head count used for every + # forward reshape / RoPE / pad / run_attention; the parallel linears are + # given the TOTAL width and divide by tp_size internally, and their + # per-rank ``weight_loader`` slices this rank's head block — so one weight + # path serves tp=1 and tp>1. The paged cache reports the matching per-rank + # head count: ``KVCacheConfig.shard`` divides ``num_qo/kv_heads`` by the + # node's instance world size (tp*sp), exactly like the Orpheus TP path. + self.tp_size = comm_group.world_size + self.total_num_heads = config.num_attention_heads + if self.total_num_heads % self.tp_size != 0: + raise ValueError( + f"num_attention_heads={self.total_num_heads} is not divisible by " + f"tp_size={self.tp_size}" + ) + self.num_heads = self.total_num_heads // self.tp_size self.qk_nope_head_dim = config.qk_nope_head_dim self.qk_rope_head_dim = config.qk_rope_head_dim self.qk_head_dim = config.qk_head_dim @@ -61,7 +83,9 @@ def __init__(self, config: KimiK2Config, comm_group: CommGroup | None = None) -> # zero-padded to this width for the paged run_attention (M6 mitigation); # the attention output is sliced back to v_head_dim. See config docstring. self.padded_head_dim = config.padded_head_dim - h = self.num_heads + # Parallel linears take the TOTAL head width (they divide by tp_size); + # the forward uses ``self.num_heads`` (local). + h = self.total_num_heads # Q: two-stage low-rank (q_a down -> norm -> q_b up). Down-projections are # replicated (small rank); up-projections shard over heads under TP. diff --git a/mstar/model/kimi_k2_7/components/moe.py b/mstar/model/kimi_k2_7/components/moe.py index 763909c11..a72f68740 100644 --- a/mstar/model/kimi_k2_7/components/moe.py +++ b/mstar/model/kimi_k2_7/components/moe.py @@ -26,6 +26,7 @@ from torch import nn from mstar.distributed.communication import CommGroup +from mstar.distributed.utils import divide from mstar.model.components.distributed import ParallelGatedMLP from mstar.model.components.moe import ( _dispatch, @@ -158,9 +159,26 @@ class KimiSparseMoeBlock(nn.Module): ``shared`` is a plain dense SwiGLU MLP added ungated (no sigmoid gate, no ``routed_scaling_factor``). - Expert weights use the fused layout reused from ``model.components.moe``: - - ``experts.gate_up_proj``: ``(E, 2 * moe_intermediate_size, hidden)`` - - ``experts.down_proj``: ``(E, hidden, moe_intermediate_size)`` + **TP sharding (intermediate-parallel).** Under tensor parallelism the router + (:class:`KimiMoEGate`) stays REPLICATED — every rank computes the full + ``(top_k_ids, weights)`` — and only the expert GEMMs shard, exactly like + mstar's own ``ParallelSparseMoeBlock``: each rank holds every expert but only + a ``moe_intermediate_size // tp_size`` slice of its SwiGLU intermediate + (``gate_up_proj`` column-parallel, ``down_proj`` row-parallel). The per-rank + partial hidden contributions are summed with a single all-reduce before the + top-k sum-reduce. The shared expert is a ``ParallelGatedMLP`` on the same comm + group, so it shards its intermediate and all-reduces internally. This reuses + the existing fused-expert machinery verbatim and is trivially goldenable + (tp>1 == tp=1). Its tradeoff: every rank still stores ALL experts' weights, so + it does NOT reduce per-rank expert memory — the 1T fit needs true + token-dispatch expert parallelism (all-to-all), which is a Phase-4 concern and + is deliberately not built here. + + Expert weights use the fused layout reused from ``model.components.moe``, + sharded to this rank (``full == moe_intermediate_size``, + ``shard == full // tp_size``): + - ``experts.gate_up_proj``: ``(E, 2 * shard, hidden)`` + - ``experts.down_proj``: ``(E, hidden, shard)`` """ def __init__( @@ -170,9 +188,13 @@ def __init__( if comm_group is None: comm_group = CommGroup.trivial() self.comm_group = comm_group + self.tp_size = comm_group.world_size + self.tp_rank = comm_group.rank self.hidden_size = config.hidden_size self.num_experts = config.n_routed_experts self.moe_intermediate_size = config.moe_intermediate_size + # Per-rank slice of each expert's SwiGLU intermediate (== full at tp=1). + shard_inter = divide(config.moe_intermediate_size, self.tp_size) self.gate = KimiMoEGate( hidden_size=config.hidden_size, @@ -190,7 +212,7 @@ def __init__( self.experts.gate_up_proj = nn.Parameter( torch.empty( config.n_routed_experts, - 2 * config.moe_intermediate_size, + 2 * shard_inter, config.hidden_size, ) ) @@ -198,15 +220,16 @@ def __init__( torch.empty( config.n_routed_experts, config.hidden_size, - config.moe_intermediate_size, + shard_inter, ) ) # The fused expert params are plain nn.Parameters, so they carry no # per-shard ``weight_loader`` by default. The M5 stacked-param rules route # each checkpoint expert via a ``"gate:N"/"up:N"/"down:N"`` shard id, so we # attach the same fused-expert loaders ``ParallelSparseMoeBlock`` uses. - # Experts are held full-size here (no expert/TP sharding yet — TODO(M6)), - # hence ``tp_rank=0, tp_size=1`` and ``full_inter == moe_intermediate_size``. + # The loaders take ``(tp_rank, tp_size, full_inter)`` and slice this rank's + # intermediate stripe out of the full-size checkpoint expert — so a single + # weight path serves tp=1 (full) and tp>1 (sharded). self._attach_expert_weight_loaders() # Ungated shared expert: a dense SwiGLU MLP with the shared intermediate @@ -230,10 +253,10 @@ def _attach_expert_weight_loaders(self) -> None: from functools import partial self.experts.gate_up_proj.weight_loader = partial( - _gate_up_weight_loader, 0, 1, self.moe_intermediate_size, + _gate_up_weight_loader, self.tp_rank, self.tp_size, self.moe_intermediate_size, ) self.experts.down_proj.weight_loader = partial( - _down_proj_weight_loader, 0, 1, self.moe_intermediate_size, + _down_proj_weight_loader, self.tp_rank, self.tp_size, self.moe_intermediate_size, ) def _apply(self, fn, recurse=True): @@ -245,14 +268,53 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: input_shape = hidden_states.shape flat = hidden_states.view(-1, self.hidden_size).contiguous() + # Router is replicated: every rank computes the full top-k selection. topk_weights, topk_ids = self.gate(flat) - routed = _dispatch( + topk_weights = topk_weights.to(flat.dtype) + if self.tp_size == 1: + routed = _dispatch( + flat, + self.experts.gate_up_proj, + self.experts.down_proj, + self.num_experts, + topk_ids, + topk_weights, + ) + else: + routed = self._dispatch_tp(flat, topk_weights, topk_ids) + # Shared expert is a ParallelGatedMLP on the same comm group: at tp>1 it + # holds its own intermediate stripe and all-reduces inside its down_proj. + shared = self.shared_expert(flat) + return (routed + shared).view(input_shape) + + def _dispatch_tp( + self, + flat: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + ) -> torch.Tensor: + """Intermediate-sharded routed dispatch (mirrors + ``ParallelSparseMoeBlock._dispatch_tp``). + + Each rank's ``fused_experts`` produces its partial hidden contribution per + (token, top-k slot); an all-reduce sums the intermediate-dim partials + across ranks, then ``moe_sum_reduce_triton`` folds the top-k dim. The + combine weights already carry ``routed_scaling_factor`` (folded in by + :class:`KimiMoEGate`), so the sum-reduce passes ``routed_scaling_factor=1.0``. + """ + from mstar.utils.fused_moe import fused_experts, moe_sum_reduce_triton + + # (tokens, top_k, hidden) partials — reduce_results=False keeps the + # per-slot rows so we can all-reduce the intermediate-parallel partials. + cache3 = fused_experts( flat, self.experts.gate_up_proj, self.experts.down_proj, - self.num_experts, + topk_weights, topk_ids, - topk_weights.to(flat.dtype), + reduce_results=False, ) - shared = self.shared_expert(flat) - return (routed + shared).view(input_shape) + self.comm_group.all_reduce(cache3) + output = torch.empty_like(flat) + moe_sum_reduce_triton(cache3, output, routed_scaling_factor=1.0) + return output diff --git a/mstar/model/kimi_k2_7/config.py b/mstar/model/kimi_k2_7/config.py index 12267cffe..798f0f9a5 100644 --- a/mstar/model/kimi_k2_7/config.py +++ b/mstar/model/kimi_k2_7/config.py @@ -87,6 +87,14 @@ class KimiK2Config: # -- MTP (multi-token prediction) — deferred, declared for completeness - num_nextn_predict_layers: int = 0 + # -- Serving: CUDA-graph prefill capture grid (optional overrides) ------ + # ``None`` => ``KimiLLMSubmodule`` uses its full-size class-default grid. + # ``reduced()`` sets a tiny grid so the synthetic bring-up serve captures a + # single short-prompt graph instead of the full 6x5 compiled grid (this was + # an env knob during bring-up; it now lives in the config). + prefill_token_buckets: list[int] | None = None + prefill_capture_batch_sizes: list[int] | None = None + # --------------------------------------------------------------------- # Derived dims (read by get_kv_cache_config / attention) # --------------------------------------------------------------------- @@ -148,4 +156,10 @@ def reduced(cls) -> "KimiK2Config": n_group=1, topk_group=1, first_k_dense_replace=1, + # Tiny CUDA-graph prefill capture grid for the synthetic bring-up serve: + # one short-prompt bucket at batch size 1 (the full 6x5 grid is slow and + # its larger buckets exceed this 512-token model). Serve/CUDA-graph path + # only — the golden tests call forward() directly and are unaffected. + prefill_token_buckets=[64], + prefill_capture_batch_sizes=[1], ) diff --git a/mstar/model/kimi_k2_7/kimi_model.py b/mstar/model/kimi_k2_7/kimi_model.py index 3149a455a..9978d300c 100644 --- a/mstar/model/kimi_k2_7/kimi_model.py +++ b/mstar/model/kimi_k2_7/kimi_model.py @@ -70,8 +70,26 @@ def __init__( **kwargs, ): self.cache_dir = cache_dir - self.model_path_hf = model_path_hf - self.config = KimiK2Config() + # ``model_kwargs`` from the serving YAML arrive here as ``**kwargs`` (see + # api_server/entrypoint.py). They let a config redirect this model at a + # local (reduced/synthetic) checkpoint without touching the shared + # registry, so a runnable text serve is possible before the 1T weights + # exist. All three are optional and default to the full-size behaviour. + # * ``checkpoint_path`` — local HF-format dir/file to load instead of + # the ``HF_MODELS`` repo id (used as-is by ``_resolve_checkpoint``). + # * ``config_variant`` — ``"reduced"`` selects ``KimiK2Config.reduced()`` + # (tiny, GPU-runnable shape); anything else keeps the 1T config. + # * ``tokenizer_mode`` — ``"byte"`` swaps the HF tokenizer for a trivial + # UTF-8 byte identity tokenizer, the pragmatic fit for the reduced + # ``vocab_size=256`` model (the real Kimi tokenizer emits ids ≫ 256). + checkpoint_path = kwargs.get("checkpoint_path") + self.model_path_hf = checkpoint_path or model_path_hf + self._config_variant = kwargs.get("config_variant", "full") + if self._config_variant == "reduced": + self.config = KimiK2Config.reduced() + else: + self.config = KimiK2Config() + self._tokenizer_mode = kwargs.get("tokenizer_mode", "hf") # Tokenizer is loaded lazily: the modular (dummy-mode) tests build the # model via ``object.__new__`` and never call ``__init__``, so we avoid # forcing a network/tokenizer dependency into the scaffold path. @@ -262,6 +280,15 @@ def process_prompt( # Text-only for M0; raw multimodal tensors (MoonViT) are a later milestone. if prompt is None: return {} + if self._tokenizer_mode == "byte": + # Trivial UTF-8 byte identity tokenizer for the reduced vocab_size=256 + # serve: each prompt byte is already a valid token id in [0, 256), so + # no HF tokenizer / network dependency is needed. Clamped defensively + # in case a smaller reduced vocab is ever used. + vocab = self.config.vocab_size + byte_ids = [min(b, vocab - 1) for b in prompt.encode("utf-8")] or [0] + input_ids = torch.tensor(byte_ids, dtype=torch.long) + return {"text_inputs": [input_ids]} input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids[0] return {"text_inputs": [input_ids]} @@ -290,6 +317,13 @@ def postprocess( ) -> bytes: if modality == "text": token_ids = output.tolist() if output.numel() else [] + if self._tokenizer_mode == "byte": + # Inverse of the byte identity tokenizer: reduced-vocab ids map + # straight back to raw bytes. The synthetic model emits arbitrary + # ids in [0, 256), so the bytes are not guaranteed valid UTF-8 — + # decode leniently (the point is to prove tokens stream, not to + # produce meaningful text on random weights). + return bytes((t & 0xFF) for t in token_ids) text = self.tokenizer.decode(token_ids, skip_special_tokens=True) return text.encode("utf-8") raise ValueError(f"Unsupported modality for Kimi-K2.7: {modality!r}") diff --git a/mstar/model/kimi_k2_7/submodules.py b/mstar/model/kimi_k2_7/submodules.py index 015ce5e5a..ed6bef6c2 100644 --- a/mstar/model/kimi_k2_7/submodules.py +++ b/mstar/model/kimi_k2_7/submodules.py @@ -56,7 +56,13 @@ def __init__(self, language_model: nn.Module, config: KimiK2Config): self.lm_head = language_model.lm_head self.config = config - # -- CUDA-graph capture buckets (mirror OrpheusLLMSubmodule) ------------ + # -- CUDA-graph prefill capture grid (full-size defaults) -------------- + # Overridable per-config via ``config.prefill_token_buckets`` / + # ``config.prefill_capture_batch_sizes``: ``KimiK2Config.reduced()`` sets a + # tiny grid for the synthetic bring-up serve, while the full model leaves them + # ``None`` and uses these defaults. Capturing the full 6x5 compiled grid is + # slow, and buckets above a small model's ``max_position_embeddings`` do not + # fit — hence the reduced config trims it (config-driven, not env-driven). PREFILL_TOKEN_BUCKETS = [32, 64, 128, 256, 512, 1024] PREFILL_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16] @@ -88,9 +94,13 @@ def get_cuda_graph_configs( above. ``inv_freq`` is lazily cached on first (warmup) forward, so its address is stable across replay. """ + prefill_buckets = self.config.prefill_token_buckets or self.PREFILL_TOKEN_BUCKETS + prefill_batch_sizes = ( + self.config.prefill_capture_batch_sizes or self.PREFILL_CAPTURE_BATCH_SIZES + ) prefill_packed = { num_tokens: self._build_prefill_packed(num_tokens, device) - for num_tokens in self.PREFILL_TOKEN_BUCKETS + for num_tokens in prefill_buckets } return [ BasicBatchedCudaGraphConfig( @@ -110,7 +120,7 @@ def get_cuda_graph_configs( labels=[_MAIN], compile=True, causal_attention=True, - capture_batch_sizes=self.PREFILL_CAPTURE_BATCH_SIZES, + capture_batch_sizes=prefill_batch_sizes, ), ] diff --git a/test/integration/test_kimi_serve_e2e.py b/test/integration/test_kimi_serve_e2e.py new file mode 100644 index 000000000..91cfaf173 --- /dev/null +++ b/test/integration/test_kimi_serve_e2e.py @@ -0,0 +1,313 @@ +"""Phase 2 (gap 4): drive the Kimi-K2.7 SERVING path as deep as possible +in-process, beyond the submodule-level gate. + +``mstar-serve`` is a multi-process stack (API server -> conductor -> N worker +processes -> KV_CACHE engine -> decode Loop -> tokens over ZMQ). A live serve is +the ultimate proof; see ``tools/kimi_goldens/repro/RUNBOOK.md`` for the exact +launch + request commands. This test is the committed, deterministic gate that +exercises the same *model* serve surface a live serve hits, minus the inter- +process transport, so it runs in CI-style isolation on one GPU with synthetic +weights: + + 1. **The real serve entry points via the real __init__.** We build the model + through ``KimiK2Model(config_variant="reduced", checkpoint_path=..., + tokenizer_mode="byte")`` — the exact path ``api_server/entrypoint.py`` takes + from a serving YAML's ``model_kwargs`` — then use ``process_prompt`` (byte + tokenizer), ``get_submodule`` (meta -> to_empty -> M5 load), and + ``postprocess``. None of these are touched by ``test_kimi_submodule.py``, + which bypasses ``__init__`` via ``object.__new__``. + + 2. **prefill + the decode Loop with the real Sampler and check_stop.** Over a + genuine ``FlashInferCacheManager`` we run prefill (``forward`` -> logits -> + ``Sampler.sample``) then several decode steps through ``forward_batched`` + (the engine's batched decode path, which samples inside the pass and returns + per-request ``new_token``), calling the submodule's ``check_stop`` each step + — the Loop's real stop mechanism. This produces actual generated token ids + and proves the decode loop terminates on ``max_tokens`` (the reduced model's + EOS id 163586 is unreachable in ``vocab_size=256``, so ``max_tokens`` is the + only stop — exactly what a live reduced serve relies on). + +Deterministic by construction (greedy / ``temperature=0``): the sequence is +asserted stable across two fresh-cache runs so the golden never flakes. + +Run: pytest test/integration/test_kimi_serve_e2e.py -v +""" +import pytest +import torch + +from mstar.communication.tensors import LocalTransferEngine +from mstar.conductor.request_info import CurrentForwardPassInfo +from mstar.engine.cache_manager import WorkspaceBufferManager, create_cache_manager +from mstar.engine.kv_store import ( + KVCacheConfig, + PagedAllocationManager, + TransferEngineInfo, +) +from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM +from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock +from mstar.model.kimi_k2_7.config import KimiK2Config +from mstar.model.kimi_k2_7.kimi_model import KimiK2Model +from mstar.model.kimi_k2_7.submodules import KimiLLMSubmodule +from mstar.model.submodule_base import ModelInputsFromEngine +from mstar.utils.sampling import Sampler, SamplingConfig + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="serve e2e needs a GPU (real FlashInfer paged cache)", +) + +DEVICE = torch.device("cuda") + + +# -------------------------------------------------------------------------- +# Synthetic HF DeepSeek-V3 reduced checkpoint (un-fuse every fused param back +# to HF keys — identical serialization to test_kimi_submodule / M5 loader). +# -------------------------------------------------------------------------- + +def _fill_layer(layer, cfg): + a = layer.self_attn + for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): + lin.weight.data.normal_(0, 0.03) + for norm in (a.q_a_layernorm, a.kv_a_layernorm): + norm.weight.data.normal_(1.0, 0.02) + layer.input_layernorm.weight.data.normal_(1.0, 0.02) + layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02) + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + mlp.gate.weight.data.normal_(0, 1) + mlp.gate.e_score_correction_bias.data = torch.randn( + cfg.n_routed_experts, device=DEVICE, dtype=torch.float32) + mlp.experts.gate_up_proj.data.normal_(0, 0.05) + mlp.experts.down_proj.data.normal_(0, 0.05) + mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05) + else: + mlp.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.down_proj.weight.data.normal_(0, 0.05) + + +def _build_reference(cfg): + model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=torch.bfloat16) + model.model.embed_tokens.weight.data.normal_(0, 0.05) + model.model.norm.weight.data.normal_(1.0, 0.02) + model.lm_head.weight.data.normal_(0, 0.02) + for layer in model.model.layers: + _fill_layer(layer, cfg) + return model.eval() + + +def _hf_checkpoint(model, cfg): + inter = cfg.intermediate_size + moe_inter = cfg.moe_intermediate_size + shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts + m = model.model + sd = {"model.embed_tokens.weight": m.embed_tokens.weight} + for i, layer in enumerate(m.layers): + p = f"model.layers.{i}." + a = layer.self_attn + sd[p + "self_attn.q_a_proj.weight"] = a.q_a_proj.weight + sd[p + "self_attn.q_a_layernorm.weight"] = a.q_a_layernorm.weight + sd[p + "self_attn.q_b_proj.weight"] = a.q_b_proj.weight + sd[p + "self_attn.kv_a_proj_with_mqa.weight"] = a.kv_a_proj_with_mqa.weight + sd[p + "self_attn.kv_a_layernorm.weight"] = a.kv_a_layernorm.weight + sd[p + "self_attn.kv_b_proj.weight"] = a.kv_b_proj.weight + sd[p + "self_attn.o_proj.weight"] = a.o_proj.weight + sd[p + "input_layernorm.weight"] = layer.input_layernorm.weight + sd[p + "post_attention_layernorm.weight"] = layer.post_attention_layernorm.weight + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + sd[p + "mlp.gate.weight"] = mlp.gate.weight + sd[p + "mlp.gate.e_score_correction_bias"] = mlp.gate.e_score_correction_bias + gup, dwn = mlp.experts.gate_up_proj, mlp.experts.down_proj + for e in range(cfg.n_routed_experts): + sd[p + f"mlp.experts.{e}.gate_proj.weight"] = gup[e, :moe_inter, :] + sd[p + f"mlp.experts.{e}.up_proj.weight"] = gup[e, moe_inter:, :] + sd[p + f"mlp.experts.{e}.down_proj.weight"] = dwn[e] + sh = mlp.shared_expert + sd[p + "mlp.shared_experts.gate_proj.weight"] = sh.gate_up_proj.weight[:shared_inter] + sd[p + "mlp.shared_experts.up_proj.weight"] = sh.gate_up_proj.weight[shared_inter:] + sd[p + "mlp.shared_experts.down_proj.weight"] = sh.down_proj.weight + else: + sd[p + "mlp.gate_proj.weight"] = mlp.gate_up_proj.weight[:inter] + sd[p + "mlp.up_proj.weight"] = mlp.gate_up_proj.weight[inter:] + sd[p + "mlp.down_proj.weight"] = mlp.down_proj.weight + sd["model.norm.weight"] = m.norm.weight + sd["lm_head.weight"] = model.lm_head.weight + return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()} + + +def _write_checkpoint(tmp_path, seed=0): + from safetensors.torch import save_file + + torch.manual_seed(seed) + cfg = KimiK2Config.reduced() + ref = _build_reference(cfg) + save_file(_hf_checkpoint(ref, cfg), str(tmp_path / "model.safetensors")) + return cfg + + +# -------------------------------------------------------------------------- +# Real paged FlashInfer cache (head_dim = padded_head_dim = 64 for reduced). +# -------------------------------------------------------------------------- + +def _make_real_cache_manager(cfg, dtype, page_size=128, max_num_pages=8): + num_heads = cfg.num_attention_heads + head_dim = cfg.padded_head_dim + kv_cache = torch.zeros( + cfg.num_hidden_layers, max_num_pages, 2, page_size, num_heads, head_dim, + dtype=dtype, device=DEVICE, + ).contiguous() + kv_cfg = KVCacheConfig( + num_layers=cfg.num_hidden_layers, num_kv_heads=num_heads, head_dim=head_dim, + max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages, + page_size=page_size, num_qo_heads=num_heads, + ) + transfer_info = TransferEngineInfo( + my_entity_id="kimi_serve_e2e", my_session_id="kimi_serve_e2e", + transfer_engine=LocalTransferEngine("localhost"), + ) + alloc = PagedAllocationManager( + config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info) + alloc.add_request("r0", ["main"]) + buffers = WorkspaceBufferManager(64 * 1024 * 1024, device=DEVICE) + cm = create_cache_manager( + request_ids=["r0"], active_labels_per_request={"r0": "main"}, + kv_cache=kv_cache, alloc_manager=alloc, buffer_manager=buffers, + kv_cache_config=kv_cfg, device=DEVICE, + ) + return cm, alloc + + +def _greedy_sampler(cfg): + """A real Sampler configured for deterministic greedy decode of r0.""" + sampler = Sampler(device=DEVICE) + sampler.add_request("r0") + sampler.set_config("r0", vocab_size=cfg.vocab_size, temperature=0.0, + top_k=0, top_p=1.0, repetition_penalty=1.0) + return sampler + + +def _fwd_info(max_tokens, cfg): + """CurrentForwardPassInfo the submodule's check_stop reads (sampling_config / + max_tokens / dynamic_loop_iter_counts).""" + return CurrentForwardPassInfo( + request_id="r0", graph_walk="decode", requires_cfg=False, fwd_index=0, + random_seed=0, max_tokens=max_tokens, + sampling_config={"LLM": SamplingConfig(vocab_size=cfg.vocab_size, + ignore_eos=cfg.ignore_eos)}, + dynamic_loop_iter_counts={}, + ) + + +# -------------------------------------------------------------------------- +# The serve drive. +# -------------------------------------------------------------------------- + +def _run_generation(model, submodule, cfg, prompt_ids, max_tokens): + """prefill (forward+Sampler) then the decode Loop (forward_batched + check_stop) + over a fresh paged cache. Returns (generated_token_ids, stopped_by_check_stop).""" + cm, alloc = _make_real_cache_manager(cfg, torch.bfloat16) + sampler = _greedy_sampler(cfg) + engine_inputs = ModelInputsFromEngine( + request_ids=["r0"], per_request_info={}, cache_manager=cm, sampler=sampler, + ) + info = _fwd_info(max_tokens, cfg) + generated: list[int] = [] + stopped = False + try: + # --- prefill: forward -> last-token logits -> Sampler (first token) --- + ar = submodule.prepare_inputs("prefill", None, {"text_inputs": [prompt_ids]}) + packed = submodule.preprocess("prefill", engine_inputs, [ar]) + with torch.no_grad(): + logits = submodule.forward("prefill", engine_inputs, **packed)["logits"][0] + assert logits.shape == (1, cfg.vocab_size) + assert torch.isfinite(logits).all() + next_token = sampler.sample(["r0"], logits).clone() # (1,) + generated.append(int(next_token.item())) + + # --- decode Loop: forward_batched samples inside the pass; check_stop --- + for step in range(max_tokens + 4): # +slack; check_stop must break first + ar = submodule.prepare_inputs("decode", None, {"text_inputs": [next_token]}) + packed = submodule.preprocess("decode", engine_inputs, [ar]) + with torch.no_grad(): + out = submodule.forward_batched("decode", engine_inputs, **packed) + new_token = out["r0"]["new_token"][0] + outputs = {"new_token": [new_token]} + submodule.postprocess("r0", info, outputs) # rebinds text_inputs + assert outputs["text_inputs"] is outputs["new_token"] + info.dynamic_loop_iter_counts["decode_loop"] = step + stop = submodule.check_stop("r0", info, outputs) + generated.append(int(new_token.item())) + next_token = new_token + if stop: + stopped = True + break + finally: + alloc.cleanup() + sampler.remove_request("r0") + return generated, stopped + + +# -------------------------------------------------------------------------- +# Tests +# -------------------------------------------------------------------------- + +def test_serve_path_prefill_decode_loop(tmp_path): + """Full model serve surface: real __init__ (reduced/local/byte) -> process_prompt + -> get_submodule -> prefill + decode Loop (Sampler + check_stop) -> postprocess. + Proves the decode loop generates real tokens and terminates on max_tokens.""" + cfg = _write_checkpoint(tmp_path, seed=0) + + # The exact construction api_server/entrypoint.py performs from a serving + # YAML's model_kwargs (no HF tokenizer, no 1T weights). + model = KimiK2Model( + model_path_hf="", config_variant="reduced", + checkpoint_path=str(tmp_path), tokenizer_mode="byte", + ) + assert model.config.vocab_size == 256 + + # Byte tokenizer: prompt text -> token ids in [0, 256). + prompt_tensors = model.process_prompt("hello kimi", ["text"], ["text"]) + prompt_ids = prompt_tensors["text_inputs"][0].to(DEVICE) + assert prompt_ids.tolist() == list("hello kimi".encode("utf-8")) + assert prompt_ids.max().item() < cfg.vocab_size + + submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16) + assert isinstance(submodule, KimiLLMSubmodule) + # M6 buffer audit still holds through the serve build path. + assert list(submodule.language_model.named_buffers()) == [] + + MAX_TOKENS = 6 + generated, stopped = _run_generation(model, submodule, cfg, prompt_ids, MAX_TOKENS) + + # The decode loop stopped via check_stop (max_tokens), NOT the safety slack. + assert stopped, "decode loop did not terminate via check_stop" + # 1 prefill token + exactly MAX_TOKENS decode tokens (check_stop fires when + # decode_loop count+1 >= max_tokens, i.e. after the MAX_TOKENS-th decode step). + assert len(generated) == 1 + MAX_TOKENS, generated + assert all(0 <= t < cfg.vocab_size for t in generated), generated + + # postprocess decodes the generated ids back to bytes (client-facing output). + out_bytes = model.postprocess(torch.tensor(generated), "text") + assert isinstance(out_bytes, bytes) + assert len(out_bytes) == len(generated) + + +def test_serve_path_is_deterministic(tmp_path): + """Greedy decode over a fresh cache is bit-stable: two runs of the whole + prefill+decode serve drive yield identical token sequences. Guards the golden + against the flakiness that has bitten this project before.""" + cfg = _write_checkpoint(tmp_path, seed=1) + model = KimiK2Model( + model_path_hf="", config_variant="reduced", + checkpoint_path=str(tmp_path), tokenizer_mode="byte", + ) + submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16) + prompt_ids = model.process_prompt("serve", ["text"], ["text"])["text_inputs"][0].to(DEVICE) + + runs = [ + _run_generation(model, submodule, cfg, prompt_ids, max_tokens=5)[0] + for _ in range(2) + ] + assert runs[0] == runs[1], runs + assert len(runs[0]) == 1 + 5 diff --git a/test/integration/test_kimi_tp.py b/test/integration/test_kimi_tp.py new file mode 100644 index 000000000..8ad47a1f0 --- /dev/null +++ b/test/integration/test_kimi_tp.py @@ -0,0 +1,351 @@ +"""Phase-3 tensor-parallel goldens for Kimi-K2.7 (reduced config): tp=2 == tp=1. + +Proves the two TP subsystems added in Phase 3 are numerically correct on the +reduced config: + + * **MLA head-sharding** (``KimiMLAAttention``): the q/kv UP-projections shard + ColumnParallel and ``o_proj`` reduces RowParallel, so each rank materializes + only its ``num_attention_heads // tp_size`` local heads. + * **MoE intermediate-sharding** (``KimiSparseMoeBlock``): the router stays + replicated; each rank holds every expert but only a + ``moe_intermediate_size // tp_size`` stripe of the SwiGLU intermediate + (gate_up column-parallel / down row-parallel), all-reduced before the top-k + sum-reduce. The shared expert shards via its ``ParallelGatedMLP``. + +Two verification levels, both keyed to the SAME deterministic source weights and +loaded through the REAL per-rank ``weight_loader`` slicing (one weight path for +tp=1 and tp>1): + + 1. **In-process rank simulation** (``test_*_tp2_sim_matches_tp1``, always runs + on one GPU): for each of the 2 ranks, build the block with ``tp_size=2`` and + that rank's weight shard, run it with a LOCAL no-op all-reduce so each rank + returns only its partial, then SUM the two ranks' partials and assert it + equals the tp=1 result within bf16 tolerance. This rigorously validates the + shard math + weight-loading; the block outputs are pure row-parallel reduces + (no un-sharded residual inside the block), so summing partials reconstructs + the reduce exactly. The live NCCL all-reduce itself rides the shared, + already-proven comm path (same primitive Orpheus tp2 uses). + + 2. **Real multi-process NCCL** (``test_tp2_nccl_matches_tp1``, runs only when + >= 2 CUDA devices are visible): spawn 2 ranks with a real NCCL comm group, + run attention + MoE + a full decoder layer with the REAL all-reduce, and + compare each rank's full output to the tp=1 reference. This exercises the + actual collective and the decoder layer's residual wiring (which the + partial-sum simulation cannot cover on its own). + +Determinism: all weights/inputs come from a fixed-seed CPU generator, so the max +abs diffs are stable run-to-run (no flakiness). + +Run: pytest test/integration/test_kimi_tp.py -v +""" +from __future__ import annotations + +import os +import socket + +import pytest +import torch + +from mstar.distributed.communication import CommGroup +from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention +from mstar.model.kimi_k2_7.components.decoder_layer import KimiDecoderLayer +from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock +from mstar.model.kimi_k2_7.config import KimiK2Config + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Kimi TP goldens need a GPU (fused expert GEMM + MLA RMSNorm are CUDA-only)", +) + +DEVICE = "cuda" +DTYPE = torch.bfloat16 +TP = 2 + + +# --------------------------------------------------------------------------- +# A world-size-2 comm group whose collectives are LOCAL no-ops. Used for the +# in-process rank simulation: each rank computes only its partial and the test +# sums the two ranks' partials to reconstruct the row-parallel all-reduce. +# --------------------------------------------------------------------------- +class _NoCommGroup(CommGroup): + def __init__(self, rank: int) -> None: + super().__init__(my_global_rank=rank, my_group_rank=rank, group_members=[0, 1]) + + def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: # noqa: D401 + return input_ + + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: + return input_ + + +class _MockMLACache: + """Paged-cache stand-in: causal SDPA at 1/sqrt(head_dim) over whatever local + heads the rank hands it (attention is per-head independent, so a rank running + SDPA on its head slice yields exactly those heads' outputs).""" + + def __init__(self, head_dim: int) -> None: + self.scale = head_dim ** -0.5 + + def set_layer_idx(self, _i): + pass + + def set_active_label(self, _l): + pass + + def advance_seq_lens(self, *_a, **_k): + pass + + def run_attention(self, q, k, v): + qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) # (H,T,D) + scores = torch.einsum("hqd,hkd->hqk", qt, kt) * self.scale + num_tokens = q.shape[0] + causal = torch.triu( + torch.full((num_tokens, num_tokens), float("-inf"), device=q.device), + diagonal=1, + ) + attn = (scores + causal).softmax(-1) + return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) + + +# --------------------------------------------------------------------------- +# Deterministic full-size source weights (CPU generator -> device-independent, +# identical across ranks / processes) and the REAL per-rank load helpers. +# --------------------------------------------------------------------------- +def _source_weights(cfg: KimiK2Config, seed: int) -> dict: + g = torch.Generator().manual_seed(seed) + H = cfg.num_attention_heads + E, Hd, I = cfg.n_routed_experts, cfg.hidden_size, cfg.moe_intermediate_size + sh = I * cfg.n_shared_experts + + def rn(*shape, std=0.03, mean=0.0): + return torch.randn(*shape, generator=g) * std + mean + + return { + # attention (replicated down-projs + norms, sharded up-projs + o_proj) + "q_a": rn(cfg.q_lora_rank, cfg.hidden_size), + "q_a_norm": rn(cfg.q_lora_rank, std=0.02, mean=1.0), + "q_b": rn(H * cfg.qk_head_dim, cfg.q_lora_rank), + "kv_a": rn(cfg.kv_lora_rank + cfg.qk_rope_head_dim, cfg.hidden_size), + "kv_a_norm": rn(cfg.kv_lora_rank, std=0.02, mean=1.0), + "kv_b": rn(H * (cfg.qk_nope_head_dim + cfg.v_head_dim), cfg.kv_lora_rank), + "o": rn(cfg.hidden_size, H * cfg.v_head_dim), + # moe (replicated fp32 router + sharded experts/shared) + "router_w": torch.randn(E, Hd, generator=g), + "router_b": torch.randn(E, generator=g), + "gate": rn(E, I, Hd, std=0.05), + "up": rn(E, I, Hd, std=0.05), + "down": rn(E, Hd, I, std=0.05), + "sh_gate": rn(sh, Hd, std=0.05), + "sh_up": rn(sh, Hd, std=0.05), + "sh_down": rn(Hd, sh, std=0.05), + # decoder-layer norms + "in_ln": rn(Hd, std=0.02, mean=1.0), + "post_ln": rn(Hd, std=0.02, mean=1.0), + } + + +def _load_attention(attn: KimiMLAAttention, src: dict) -> None: + """Load full source weights through the REAL Column/Row weight_loaders (which + slice this rank's head block) + direct copies for the replicated params.""" + attn.q_a_proj.weight.data.copy_(src["q_a"].to(DEVICE, DTYPE)) + attn.q_a_layernorm.weight.data.copy_(src["q_a_norm"].to(DEVICE, DTYPE)) + attn.q_b_proj.weight.weight_loader(attn.q_b_proj.weight, src["q_b"].to(DEVICE, DTYPE)) + attn.kv_a_proj_with_mqa.weight.data.copy_(src["kv_a"].to(DEVICE, DTYPE)) + attn.kv_a_layernorm.weight.data.copy_(src["kv_a_norm"].to(DEVICE, DTYPE)) + attn.kv_b_proj.weight.weight_loader(attn.kv_b_proj.weight, src["kv_b"].to(DEVICE, DTYPE)) + attn.o_proj.weight.weight_loader(attn.o_proj.weight, src["o"].to(DEVICE, DTYPE)) + + +def _load_moe(block: KimiSparseMoeBlock, src: dict) -> None: + """Load through the REAL fused-expert weight_loaders (per-rank intermediate + slice) + replicated fp32 router + the shared expert's merged/row loaders.""" + block.gate.weight.data = src["router_w"].to(DEVICE) # keep router fp32 + block.gate.e_score_correction_bias.data = src["router_b"].to(DEVICE) + gu, dp = block.experts.gate_up_proj, block.experts.down_proj + for e in range(block.num_experts): + gu.weight_loader(gu, src["gate"][e].to(DEVICE, DTYPE), loaded_shard_id=f"gate:{e}") + gu.weight_loader(gu, src["up"][e].to(DEVICE, DTYPE), loaded_shard_id=f"up:{e}") + dp.weight_loader(dp, src["down"][e].to(DEVICE, DTYPE), loaded_shard_id=f"down:{e}") + s = block.shared_expert + s.gate_up_proj.weight.weight_loader( + s.gate_up_proj.weight, src["sh_gate"].to(DEVICE, DTYPE), loaded_shard_id=0) + s.gate_up_proj.weight.weight_loader( + s.gate_up_proj.weight, src["sh_up"].to(DEVICE, DTYPE), loaded_shard_id=1) + s.down_proj.weight.weight_loader(s.down_proj.weight, src["sh_down"].to(DEVICE, DTYPE)) + + +def _load_decoder(layer: KimiDecoderLayer, src: dict) -> None: + _load_attention(layer.self_attn, src) + _load_moe(layer.mlp, src) + layer.input_layernorm.weight.data.copy_(src["in_ln"].to(DEVICE, DTYPE)) + layer.post_attention_layernorm.weight.data.copy_(src["post_ln"].to(DEVICE, DTYPE)) + + +def _inputs(cfg: KimiK2Config, num_tokens: int, seed: int): + g = torch.Generator().manual_seed(seed) + h = (torch.randn(num_tokens, cfg.hidden_size, generator=g) * 0.1).to(DEVICE, DTYPE) + pos = torch.arange(num_tokens, device=DEVICE) + return h, pos + + +# --------------------------------------------------------------------------- +# Level 1 — in-process rank simulation (single GPU, deterministic, always runs) +# --------------------------------------------------------------------------- +def test_mla_attention_tp2_sim_matches_tp1(): + cfg = KimiK2Config.reduced() + src = _source_weights(cfg, seed=101) + h, pos = _inputs(cfg, num_tokens=6, seed=202) + + ref = KimiMLAAttention(cfg, CommGroup.trivial()).to(DEVICE, DTYPE) + _load_attention(ref, src) + assert ref.num_heads == cfg.num_attention_heads # tp=1 sees all heads + out_ref = ref(h, _MockMLACache(cfg.padded_head_dim), pos) + + partials = [] + for rank in range(TP): + attn = KimiMLAAttention(cfg, _NoCommGroup(rank)).to(DEVICE, DTYPE) + assert attn.num_heads == cfg.num_attention_heads // TP # rank sees local heads + _load_attention(attn, src) + partials.append(attn(h, _MockMLACache(cfg.padded_head_dim), pos)) + + out_tp2 = partials[0] + partials[1] # row-parallel o_proj reduce == sum of ranks + max_abs = (out_tp2 - out_ref).abs().max().item() + assert max_abs < 5e-2, f"MLA tp2 vs tp1 max abs diff {max_abs}" + torch.testing.assert_close(out_tp2, out_ref, rtol=2e-2, atol=2e-2) + + +def test_moe_block_tp2_sim_matches_tp1(): + cfg = KimiK2Config.reduced() + src = _source_weights(cfg, seed=303) + h, _ = _inputs(cfg, num_tokens=7, seed=404) + + ref = KimiSparseMoeBlock(cfg, CommGroup.trivial()).to(DEVICE, DTYPE) + _load_moe(ref, src) + full_inter = ref.experts.gate_up_proj.shape[1] + out_ref = ref(h) + + partials = [] + for rank in range(TP): + block = KimiSparseMoeBlock(cfg, _NoCommGroup(rank)).to(DEVICE, DTYPE) + # each rank holds only a 1/TP stripe of the fused intermediate + assert block.experts.gate_up_proj.shape[1] == full_inter // TP + _load_moe(block, src) + partials.append(block(h)) + + out_tp2 = partials[0] + partials[1] # intermediate-parallel reduce == sum of ranks + max_abs = (out_tp2 - out_ref).abs().max().item() + assert max_abs < 5e-2, f"MoE tp2 vs tp1 max abs diff {max_abs}" + torch.testing.assert_close(out_tp2, out_ref, rtol=2e-2, atol=2e-2) + + +def test_tp2_sim_is_stable_across_repeats(): + """The simulated tp2==tp1 diffs must be bit-stable across repeats (no flaky + golden). Re-run the attention + MoE simulation 3x from the same seeds and + assert an identical max abs diff each time.""" + cfg = KimiK2Config.reduced() + + def attn_diff(): + src = _source_weights(cfg, seed=101) + h, pos = _inputs(cfg, num_tokens=6, seed=202) + ref = KimiMLAAttention(cfg, CommGroup.trivial()).to(DEVICE, DTYPE) + _load_attention(ref, src) + o_ref = ref(h, _MockMLACache(cfg.padded_head_dim), pos) + parts = [] + for rank in range(TP): + a = KimiMLAAttention(cfg, _NoCommGroup(rank)).to(DEVICE, DTYPE) + _load_attention(a, src) + parts.append(a(h, _MockMLACache(cfg.padded_head_dim), pos)) + return (parts[0] + parts[1] - o_ref).abs().max().item() + + diffs = [attn_diff() for _ in range(3)] + assert diffs[0] == diffs[1] == diffs[2], f"unstable tp2 sim diffs: {diffs}" + + +# --------------------------------------------------------------------------- +# Level 2 — real multi-process NCCL (runs only with >= 2 CUDA devices) +# --------------------------------------------------------------------------- +def _free_port() -> int: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + + +def _nccl_worker(rank: int, world_size: int, port: int, result_path: str) -> None: + """One TP rank: real NCCL comm group, run attention + MoE + decoder layer, and + compare each to a tp=1 reference built in-process from the same source weights. + Rank 0 writes the max abs diffs to ``result_path``.""" + import torch.distributed as dist + + os.environ.setdefault("NCCL_IB_DISABLE", "1") # coriander has no RDMA/IB + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=f"tcp://127.0.0.1:{port}", + world_size=world_size, + rank=rank, + ) + try: + cfg = KimiK2Config.reduced() + src = _source_weights(cfg, seed=505) + h, pos = _inputs(cfg, num_tokens=6, seed=606) + + cg = CommGroup(my_global_rank=rank, my_group_rank=rank, group_members=[0, 1]) + cg.device_group = None # 2-rank world == default group + cg.initialized = True + + # tp=2 blocks with the REAL all-reduce -> each rank produces the FULL output. + attn = KimiMLAAttention(cfg, cg).to(DEVICE, DTYPE) + _load_attention(attn, src) + moe = KimiSparseMoeBlock(cfg, cg).to(DEVICE, DTYPE) + _load_moe(moe, src) + dec = KimiDecoderLayer(cfg, layer_idx=1, comm_group=cg).to(DEVICE, DTYPE) + _load_decoder(dec, src) + assert isinstance(dec.mlp, KimiSparseMoeBlock) # layer_idx=1 is a MoE layer + + attn_tp2 = attn(h, _MockMLACache(cfg.padded_head_dim), pos) + moe_tp2 = moe(h) + dec_tp2 = dec(h, _MockMLACache(cfg.padded_head_dim), pos) + + # tp=1 reference (trivial group, same source weights). + attn_ref = KimiMLAAttention(cfg, CommGroup.trivial()).to(DEVICE, DTYPE) + _load_attention(attn_ref, src) + moe_ref = KimiSparseMoeBlock(cfg, CommGroup.trivial()).to(DEVICE, DTYPE) + _load_moe(moe_ref, src) + dec_ref = KimiDecoderLayer(cfg, layer_idx=1, comm_group=CommGroup.trivial()).to(DEVICE, DTYPE) + _load_decoder(dec_ref, src) + + o_attn = attn_ref(h, _MockMLACache(cfg.padded_head_dim), pos) + o_moe = moe_ref(h) + o_dec = dec_ref(h, _MockMLACache(cfg.padded_head_dim), pos) + + diffs = { + "attn": (attn_tp2 - o_attn).abs().max().item(), + "moe": (moe_tp2 - o_moe).abs().max().item(), + "decoder": (dec_tp2 - o_dec).abs().max().item(), + } + dist.barrier() + if rank == 0: + torch.save(diffs, result_path) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + torch.cuda.device_count() < 2, + reason="real NCCL tp=2 golden needs >= 2 CUDA devices", +) +def test_tp2_nccl_matches_tp1(tmp_path): + import torch.multiprocessing as mp + + result_path = str(tmp_path / "tp2_diffs.pt") + port = _free_port() + mp.spawn(_nccl_worker, args=(TP, port, result_path), nprocs=TP, join=True) + + diffs = torch.load(result_path) + # Real all-reduce -> each rank's full output must match the tp=1 reference + # within bf16 tolerance (attention + MoE + decoder-layer residual wiring). + assert diffs["attn"] < 5e-2, diffs + assert diffs["moe"] < 5e-2, diffs + assert diffs["decoder"] < 5e-2, diffs From a84883afab667780aabd861def58447a07359bf4 Mon Sep 17 00:00:00 2001 From: Garv Ghai Date: Sat, 25 Jul 2026 20:40:20 +0000 Subject: [PATCH 3/9] Kimi-K2.7-Code: INT4 W4A16 serving + real-checkpoint wiring - quantization.py: compressed-tensors INT4 parser (dequant-on-load) + Triton fused_moe_kernel_w4a16 (in-kernel dequant of packed int32 experts). KimiSparseMoeBlock keeps experts packed when quantized; bf16 fused-MoE path (shared with Qwen3-Omni) unchanged. - config/kimi_model/weight_loader: k27_code config, nested text_config.quantization_config auto-read, language_model. prefix strip, packed-expert stacked rules; TP8 serve configs (kimi_k2_7_code_tp8[_shm].yaml, SHM = proven load path). - TP robustness: 2h NCCL timeout for large-checkpoint bring-up; worker tp-leader gate on new-token counting. Tests: CPU quant/wiring goldens + GPU quant/kernel goldens. --- configs/kimi_k2_7_code_tp8.yaml | 47 +++ configs/kimi_k2_7_code_tp8_shm.yaml | 20 + configs/kimi_k2_7_repro.yaml | 4 +- configs/kimi_k2_7_tp2.yaml | 4 +- mstar/distributed/communication.py | 11 +- mstar/model/kimi_k2_7/_testing.py | 63 +++ mstar/model/kimi_k2_7/components/attention.py | 14 +- mstar/model/kimi_k2_7/components/causal_lm.py | 12 +- .../kimi_k2_7/components/decoder_layer.py | 2 +- .../kimi_k2_7/components/language_model.py | 32 +- mstar/model/kimi_k2_7/components/moe.py | 227 +++++++++-- mstar/model/kimi_k2_7/components/rope.py | 6 +- mstar/model/kimi_k2_7/config.py | 127 ++++-- mstar/model/kimi_k2_7/kimi_model.py | 92 ++++- mstar/model/kimi_k2_7/quantization.py | 245 ++++++++++++ mstar/model/kimi_k2_7/submodules.py | 10 +- mstar/model/kimi_k2_7/weight_loader.py | 264 ++++++------ mstar/utils/fused_moe/kernels.py | 268 ++++++++++++- mstar/utils/fused_moe/runner.py | 162 ++++++-- mstar/worker/worker.py | 9 +- .../test_kimi_moe_inkernel_dequant.py | 123 ++++++ ...test_kimi_quant_inkernel_weight_loading.py | 376 ++++++++++++++++++ .../test_kimi_quant_weight_loading.py | 269 +++++++++++++ test/modular/test_kimi_k27_code_wiring.py | 220 ++++++++++ test/modular/test_kimi_quant.py | 278 +++++++++++++ 25 files changed, 2603 insertions(+), 282 deletions(-) create mode 100644 configs/kimi_k2_7_code_tp8.yaml create mode 100644 configs/kimi_k2_7_code_tp8_shm.yaml create mode 100644 mstar/model/kimi_k2_7/_testing.py create mode 100644 mstar/model/kimi_k2_7/quantization.py create mode 100644 test/integration/test_kimi_moe_inkernel_dequant.py create mode 100644 test/integration/test_kimi_quant_inkernel_weight_loading.py create mode 100644 test/integration/test_kimi_quant_weight_loading.py create mode 100644 test/modular/test_kimi_k27_code_wiring.py create mode 100644 test/modular/test_kimi_quant.py diff --git a/configs/kimi_k2_7_code_tp8.yaml b/configs/kimi_k2_7_code_tp8.yaml new file mode 100644 index 000000000..35ee7376d --- /dev/null +++ b/configs/kimi_k2_7_code_tp8.yaml @@ -0,0 +1,47 @@ +model: "kimi_k2_7" +# FALLBACK ONLY — prefer kimi_k2_7_code_tp8_shm.yaml +# Kimi-K2.7 text backbone — REAL moonshotai/Kimi-K2.7-Code single-node TP=8 serve +# config (INT4, text-only). This points at the actual ~595 GB multimodal +# KimiK25ForConditionalGeneration checkpoint but serves ONLY its DeepSeek-V3 text +# path: the DeepSeek-V3 dims live under `text_config`, its `quantization_config` +# (pack-quantized, num_bits=4, group_size=32, routed experts only) is NESTED under +# `text_config`, and the vision tower (vision_tower.* / mm_projector.*) is dropped +# on load. Serves the routed experts packed with in-kernel Triton W4A16 dequant — +# the only path that fits the real weights. +# +# `model_kwargs` (forwarded to KimiK2Model.__init__): +# * config_variant: k27_code -> KimiK2Config.k27_code() (full 1T dims + +# moe_in_kernel_dequant=True; keeps the default +# beta_fast=32.0 — the K2.7-Code text_config value). +# The checkpoint's nested +# `text_config.quantization_config` is auto-read at +# load, so the routed experts are served packed (int32); +# the lm_head / MLA / dense-FFN / shared experts / +# vision stay bf16, matching the checkpoint `ignore`. +# * checkpoint_path -> local HF-format snapshot of the 595 GB download +# * tokenizer_mode: hf -> the real Kimi tokenizer (tiktoken-based; ids >> 256) +# +# max_seq_len is trimmed to 8192 for first bring-up (the real max is 262144); keep +# the paged KV cache modest for a first single-request serve. +# +# On this cluster (coriander) RDMA is unavailable, so launch with a non-RDMA tensor +# transport: TENSOR_PROTOCOL=SHM (single node, all 8 GPUs) — see mstar/.sample.env. +max_seq_len: 8192 +model_kwargs: + config_variant: k27_code + checkpoint_path: /m-coriander/coriander/garv901/kimi_k2_7_code + tokenizer_mode: hf +# Paged KV cache sized for the real MLA head dims (num_attention_heads=64 KV heads +# in the naive/materialized MLA path, padded_head_dim=256). 512 pages * 128 = +# 65536-token capacity — modest, for a first single-request bring-up serve. +kv_cache: + max_num_pages: 512 + page_size: 128 +# The LLM node is the tensor-parallel node: all 8 ranks on one node with tp_size 8 +# for both the prefill and decode walks. The model code is TP-agnostic — the +# per-rank shard degree comes from here. +node_groups: + - node_names: [LLM] + ranks: [0, 1, 2, 3, 4, 5, 6, 7] + tp_size: 8 + graph_walks: [prefill, decode] diff --git a/configs/kimi_k2_7_code_tp8_shm.yaml b/configs/kimi_k2_7_code_tp8_shm.yaml new file mode 100644 index 000000000..1b8ec258c --- /dev/null +++ b/configs/kimi_k2_7_code_tp8_shm.yaml @@ -0,0 +1,20 @@ +model: "kimi_k2_7" +# Attempt-3 variant of kimi_k2_7_code_tp8.yaml: identical EXCEPT the checkpoint is +# loaded from /dev/shm (RAM-backed tmpfs) instead of shared ZFS. Attempt 2 loaded +# from ZFS over ~59 min and a rank was silently SIGKILL'd (likely a shared-FS mmap +# SIGBUS / contention reaper during the long load). Staging the 595GB checkpoint +# into /dev/shm makes the load RAM-speed (minutes) and immune to mmap I/O faults. +# Stage first with tools: /m-coriander/coriander/garv901/stage_to_shm.sh +max_seq_len: 8192 +model_kwargs: + config_variant: k27_code + checkpoint_path: /dev/shm/kimi_k2_7_code + tokenizer_mode: hf +kv_cache: + max_num_pages: 512 + page_size: 128 +node_groups: + - node_names: [LLM] + ranks: [0, 1, 2, 3, 4, 5, 6, 7] + tp_size: 8 + graph_walks: [prefill, decode] diff --git a/configs/kimi_k2_7_repro.yaml b/configs/kimi_k2_7_repro.yaml index f284251eb..14f0d86e8 100644 --- a/configs/kimi_k2_7_repro.yaml +++ b/configs/kimi_k2_7_repro.yaml @@ -1,6 +1,6 @@ model: "kimi_k2_7" -# Kimi-K2.7 text backbone — REDUCED / synthetic single-GPU repro config (Phase 2 -# gap-4 serve bring-up). It drives the full serving path (API server -> conductor +# Kimi-K2.7 text backbone — REDUCED / synthetic single-GPU repro config (serve +# bring-up). It drives the full serving path (API server -> conductor # -> worker -> KV_CACHE engine -> decode loop -> tokens) on a tiny model that runs # without the 1T checkpoint. NOT a real deployment — for that use kimi_k2_7.yaml. # diff --git a/configs/kimi_k2_7_tp2.yaml b/configs/kimi_k2_7_tp2.yaml index 04de97a0b..1bbf87e5c 100644 --- a/configs/kimi_k2_7_tp2.yaml +++ b/configs/kimi_k2_7_tp2.yaml @@ -1,6 +1,6 @@ model: "kimi_k2_7" -# Kimi-K2.7 text backbone — REDUCED / synthetic TP=2 repro config (Phase 3 gap-3 -# TP correctness). Same tiny synthetic model as kimi_k2_7_repro.yaml, but the LLM +# Kimi-K2.7 text backbone — REDUCED / synthetic TP=2 repro config (TP +# correctness). Same tiny synthetic model as kimi_k2_7_repro.yaml, but the LLM # node runs tensor-parallel across 2 ranks so the MLA head-shard + MoE # intermediate-shard paths execute end-to-end. NOT a real deployment — the 1T # model needs TP8 / multi-node (see kimi_k2_7.yaml). diff --git a/mstar/distributed/communication.py b/mstar/distributed/communication.py index 2f2572f26..f2f7101b2 100644 --- a/mstar/distributed/communication.py +++ b/mstar/distributed/communication.py @@ -1,4 +1,5 @@ from dataclasses import dataclass, field +from datetime import timedelta from typing import Any import torch @@ -233,11 +234,15 @@ def init_dist( if not self.any_parallelism: return + # A generous timeout (default is ~10 min). First-time bring-up of a very + # large checkpoint — slow weight load, first-ever kernel JIT, and CUDA-graph + # capture can leave ranks waiting at this setup barrier past the default. dist.init_process_group( backend="nccl", init_method=init_method, world_size=self.num_workers, rank=self.global_rank, + timeout=timedelta(hours=2), ) # One subgroup per distinct rank tuple across BOTH mesh axes — @@ -247,7 +252,11 @@ def init_dist( # an SP group (degenerate meshes) maps to one subgroup. rank_tuple_to_pg: dict[tuple[int, ...], "dist.ProcessGroup"] = {} for rank_tuple in self.world_parallel_groups: - rank_tuple_to_pg[rank_tuple] = dist.new_group(ranks=list(rank_tuple)) + # Same generous timeout as init_process_group above (slow first load, + # kernel JIT, and graph capture can stall ranks past the default). + rank_tuple_to_pg[rank_tuple] = dist.new_group( + ranks=list(rank_tuple), timeout=timedelta(hours=2) + ) seen: set[int] = set() for comm_group in ( diff --git a/mstar/model/kimi_k2_7/_testing.py b/mstar/model/kimi_k2_7/_testing.py new file mode 100644 index 000000000..ae64329f3 --- /dev/null +++ b/mstar/model/kimi_k2_7/_testing.py @@ -0,0 +1,63 @@ +"""Test-support helpers for the Kimi-K2.7 compressed-tensors path. + +NOT part of the serving path. These helpers exist only to let the test suite +fabricate a synthetic quantized checkpoint and its exact bf16 reference, so a +golden can assert the real load path (``quantization.py`` + +``weight_loader.py``) reproduces the reference bit-for-bit. Nothing here is +imported by the model or the loader at serve time. + +The real load-path primitives (``pack_int32`` / ``unpack_int32`` / +``dequantize_weight`` / ``dequant_compressed_tensors_stream`` / +``CompressedTensorsQuantConfig``) live in ``quantization.py``; this module builds +on them. +""" +from __future__ import annotations + +import torch + +from mstar.model.kimi_k2_7.quantization import dequantize_weight, pack_int32 + + +def fake_quantize_weight( + weight: torch.Tensor, + *, + num_bits: int, + group_size: int, + symmetric: bool = True, + scale_dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize ``weight`` group-wise and return ``(packed, scale, dequant)``. + + A test/harness helper (not used at serve time): produces the on-disk + compressed-tensors tensors *and* the exact bf16 result they dequantize back + to, so a golden can assert the loader reproduces ``dequant`` bit-for-bit. Only + symmetric INT-style quantization is implemented (Kimi's scheme). + + ``dequant`` is derived from the *returned* ``scale`` via :func:`dequantize_weight`, + so it stays consistent with whatever ``scale_dtype`` the scale is stored at. + Pass ``scale_dtype=torch.bfloat16`` to match a real compressed-tensors + checkpoint (whose ``weight_scale`` is stored in the model dtype) — otherwise + the loader's bf16-scale dequant would differ from an fp32-scale reference in + the low bits. + """ + if not symmetric: + raise NotImplementedError("fake_quantize_weight: only symmetric is implemented") + out_f, in_f = weight.shape + gs = in_f if group_size in (-1, None) else group_size + if in_f % gs != 0: + raise ValueError(f"in_features {in_f} not divisible by group_size {gs}") + + qmax = (1 << (num_bits - 1)) - 1 # 7 for INT4 + qmin = -(1 << (num_bits - 1)) # -8 + w = weight.to(torch.float32).reshape(out_f, in_f // gs, gs) + scale = w.abs().amax(dim=-1, keepdim=True) / qmax + scale = torch.where(scale == 0, torch.ones_like(scale), scale) + q = torch.round(w / scale).clamp(qmin, qmax) # signed [qmin, qmax], fp32 scale + + q_unsigned = (q + float(1 << (num_bits - 1))).reshape(out_f, in_f).to(torch.int64) + packed = pack_int32(q_unsigned, num_bits) + scale_2d = scale.squeeze(-1).to(scale_dtype) + dequant = dequantize_weight( + packed, scale_2d, num_bits=num_bits, group_size=group_size, symmetric=symmetric, + ) + return packed, scale_2d, dequant diff --git a/mstar/model/kimi_k2_7/components/attention.py b/mstar/model/kimi_k2_7/components/attention.py index 30a63d3cf..fc861d17a 100644 --- a/mstar/model/kimi_k2_7/components/attention.py +++ b/mstar/model/kimi_k2_7/components/attention.py @@ -16,8 +16,8 @@ ``Drope`` slice is the single shared MQA rope key ``k_pe[T,1,Drope]``. - YARN RoPE rotates only ``q_pe`` (per head) and ``k_pe`` (broadcast to H heads). - assemble ``k = [k_nope | k_pe_broadcast] -> [T,H,Dqk]``; zero-pad ``q``/``k`` - (Dqk) and ``v`` (Dv) up to ``padded_head_dim`` (M6 mitigation — FlashInfer SM90 - rejects ``head_dim_vo`` not in {64,128,256}); fold the scale boost into ``q`` + (Dqk) and ``v`` (Dv) up to ``padded_head_dim`` (FlashInfer SM90 rejects + ``head_dim_vo`` not in {64,128,256}); fold the scale boost into ``q`` (``run_attention`` uses the fixed ``1/sqrt(padded_head_dim)`` scale), attend, slice the output back to ``Dv``, ``o_proj``. @@ -26,8 +26,10 @@ Dqk=24). Under tensor parallelism each rank materializes only its ``num_attention_heads // tp_size`` local heads (K/V and Q both shard on the head axis — there is no separate KV-head group in the naive path), and the paged cache -reports the matching per-rank count. Weight-absorbed MLA (native latent dims, no -pad) and the ``fused_qkv_a_proj`` weight fusion are deferred to the perf backlog. +reports the matching per-rank count. + +TODO: weight-absorbed MLA (native latent dims, no pad) and the ``fused_qkv_a_proj`` +weight fusion are not implemented yet. """ from __future__ import annotations @@ -80,8 +82,8 @@ def __init__(self, config: KimiK2Config, comm_group: CommGroup | None = None) -> self.v_head_dim = config.v_head_dim self.kv_lora_rank = config.kv_lora_rank # FlashInfer SM90 rejects head_dim_vo not in {64,128,256}, so q/k/v are - # zero-padded to this width for the paged run_attention (M6 mitigation); - # the attention output is sliced back to v_head_dim. See config docstring. + # zero-padded to this width for the paged run_attention; the attention + # output is sliced back to v_head_dim. See config docstring. self.padded_head_dim = config.padded_head_dim # Parallel linears take the TOTAL head width (they divide by tp_size); # the forward uses ``self.num_heads`` (local). diff --git a/mstar/model/kimi_k2_7/components/causal_lm.py b/mstar/model/kimi_k2_7/components/causal_lm.py index 7d45f517d..8920927e9 100644 --- a/mstar/model/kimi_k2_7/components/causal_lm.py +++ b/mstar/model/kimi_k2_7/components/causal_lm.py @@ -1,6 +1,6 @@ -"""Kimi-K2.7 / DeepSeek-V3 assembled text backbone (M4 assembly). +"""Kimi-K2.7 / DeepSeek-V3 assembled text backbone. -Stacks the M4 :class:`KimiDecoderLayer` blocks between a token embedding and a +Stacks the :class:`KimiDecoderLayer` blocks between a token embedding and a final RMSNorm (:class:`KimiLanguageModel`), then wraps that with the untied LM head (:class:`KimiForCausalLM`). This is the full text forward: token ids → logits. @@ -87,7 +87,7 @@ def forward( return self.lm_head(hidden_states) def load_weights(self, weights, **kwargs) -> set[str]: - """Load an HF DeepSeek-V3 checkpoint stream (M5). + """Load an HF DeepSeek-V3 checkpoint stream. Called by the shared ``mstar.model.loader.load_weights(model, source, device)`` driver (mirrors ``OrpheusForCausalLM.load_weights``). Delegates @@ -97,6 +97,12 @@ def load_weights(self, weights, **kwargs) -> set[str]: """ from mstar.model.kimi_k2_7.weight_loader import load_kimi_hf_weights + packed_experts = ( + self.config.quantization_config is not None + and self.config.moe_in_kernel_dequant + ) return load_kimi_hf_weights( self, weights, self.config.n_routed_experts, + quant_config=self.config.quantization_config, + packed_experts=packed_experts, ) diff --git a/mstar/model/kimi_k2_7/components/decoder_layer.py b/mstar/model/kimi_k2_7/components/decoder_layer.py index 88e581d01..5048defff 100644 --- a/mstar/model/kimi_k2_7/components/decoder_layer.py +++ b/mstar/model/kimi_k2_7/components/decoder_layer.py @@ -1,4 +1,4 @@ -"""Kimi-K2.7 / DeepSeek-V3 decoder layer (M4 assembly). +"""Kimi-K2.7 / DeepSeek-V3 decoder layer. One pre-norm transformer block: MLA self-attention then a feed-forward that is either the dense SwiGLU MLP (the ``first_k_dense_replace`` early layers) or the diff --git a/mstar/model/kimi_k2_7/components/language_model.py b/mstar/model/kimi_k2_7/components/language_model.py index 6edcf5ef1..cc05ad7e1 100644 --- a/mstar/model/kimi_k2_7/components/language_model.py +++ b/mstar/model/kimi_k2_7/components/language_model.py @@ -1,21 +1,13 @@ -"""Kimi-K2.7 language-model components (DeepSeek-V3 text backbone). - -M1 (cheap reuse) wiring: builders that map ``KimiK2Config`` onto mstar's existing -reused primitives — token embedding, the dense SwiGLU MLP (the -``first_k_dense_replace`` early layers), RMSNorm, and the LM head. These are the -pieces DeepSeek-V3 shares verbatim with a standard Llama-style stack; the -Kimi-specific parts (MLA attention, fine-grained sigmoid-routed MoE, YARN RoPE) -are separate milestones (M2/M3). - -Each builder is thin on purpose: it fixes the config→component mapping (dims, -``silu`` activation, ``bias=False``, RMSNorm eps, tied-vs-untied LM head) that M4 -assembles into the full ``KimiLanguageModel``. Every builder matches the vLLM -DeepSeek-V3 reference: - - embedding / LM head: ``deepseek_v2.py`` ``DeepseekV2ForCausalLM`` (untied, - ``tie_word_embeddings=False``); - - dense MLP: ``DeepseekV2MLP`` = ``down_proj(SiluAndMul(gate_up_proj(x)))``, - ``bias=False``, silu-only; - - RMSNorm: standard Llama-style ``x * rsqrt(mean(x^2)+eps) * weight``. +"""Kimi-K2.7 language-model builders (DeepSeek-V3 text backbone). + +Thin builders mapping ``KimiK2Config`` onto reused mstar primitives — token +embedding, the dense SwiGLU MLP (the ``first_k_dense_replace`` early layers), +RMSNorm, and the untied LM head — the pieces DeepSeek-V3 shares with a standard +Llama-style stack (the Kimi-specific MLA attention, sigmoid-routed MoE, and YARN +RoPE live elsewhere). Each matches the vLLM DeepSeek-V3 reference +(``deepseek_v2.py``): untied embedding/LM head (``tie_word_embeddings=False``), +``DeepseekV2MLP`` = ``down_proj(SiluAndMul(gate_up_proj(x)))`` (``bias=False``, +silu-only), Llama-style RMSNorm. """ from __future__ import annotations @@ -72,7 +64,7 @@ def build_dense_mlp( Matches ``DeepseekV2MLP``: fused gate/up projection, ``silu(gate) * up``, row-parallel down projection, ``bias=False``. Uses the full ``intermediate_size`` (the MoE layers use ``moe_intermediate_size`` per - expert instead — that path is M2). + expert instead). """ return ParallelGatedMLP( hidden_size=config.hidden_size, @@ -109,7 +101,7 @@ def build_mlp_for_layer( Returns a ``ParallelGatedMLP`` for the early dense layers, else a ``KimiSparseMoeBlock``. Both expose the same ``(x) -> x`` interface, so the - decoder layer (M4) is agnostic to which it holds. + decoder layer is agnostic to which it holds. """ if is_moe_layer(config, layer_idx): return build_moe_block(config, comm_group=comm_group) diff --git a/mstar/model/kimi_k2_7/components/moe.py b/mstar/model/kimi_k2_7/components/moe.py index a72f68740..47a26c375 100644 --- a/mstar/model/kimi_k2_7/components/moe.py +++ b/mstar/model/kimi_k2_7/components/moe.py @@ -35,6 +35,69 @@ ) from mstar.model.kimi_k2_7.config import KimiK2Config +# --------------------------------------------------------------------------- +# Packed-expert weight loaders (int32 weights + bf16 group scales). +# +# The packed analogs of ``model.components.moe._gate_up_weight_loader`` / +# ``_down_proj_weight_loader`` (which serve the bf16 fused params shared with +# Qwen3-Omni). The TP shard geometry is identical to the bf16 loaders; only the +# last (input/K) axis differs: it is pre-divided by ``pack_factor`` (packed int32) +# or ``group_size`` (bf16 scale). One function serves both the packed and the +# scale tensor for a projection — the divisor is the only difference. +# --------------------------------------------------------------------------- + + +def _gate_up_packed_loader( + tp_rank: int, tp_size: int, full_inter: int, + param: nn.Parameter, loaded_weight: torch.Tensor, + loaded_shard_id: str | int | None = None, +): + """Load one expert's gate_proj/up_proj packed-or-scale tensor into the fused + ``gate_up_proj_packed`` / ``gate_up_proj_scale`` param. + + ``loaded_shard_id`` is ``"gate:N"`` / ``"up:N"``. ``loaded_weight`` is a single + expert's 2-D tensor ``(full_inter, hidden // divisor)`` (divisor = pack_factor + for the int32 packed tensor, group_size for the bf16 scale). The N/out axis + (dim 0) is the TP-sharded one: this rank takes rows + ``[tp_rank*shard_inter : +shard_inter]`` and writes them into the gate half + ``[:shard_inter]`` or up half ``[shard_inter:]`` of ``param[expert]``. The last + axis (the un-sharded input dim) is copied whole. + """ + assert loaded_shard_id is not None + kind, expert_str = loaded_shard_id.split(":") + expert_idx = int(expert_str) + shard_inter = divide(full_inter, tp_size) + start = tp_rank * shard_inter + tp_slice = loaded_weight[start:start + shard_inter, :] + if kind == "gate": + param.data[expert_idx, :shard_inter, :] = tp_slice + else: + param.data[expert_idx, shard_inter:, :] = tp_slice + + +def _down_packed_loader( + tp_rank: int, tp_size: int, full_inter: int, divisor: int, + param: nn.Parameter, loaded_weight: torch.Tensor, + loaded_shard_id: str | int | None = None, +): + """Load one expert's down_proj packed-or-scale tensor into ``down_proj_packed`` + / ``down_proj_scale``. + + ``loaded_shard_id`` is ``"down:N"``. ``loaded_weight`` is ``(hidden, moe_inter + // divisor)``; the intermediate dim is the LAST (input) axis and is the + TP-sharded one, already divided by ``divisor`` (pack_factor for packed, + group_size for scale). This rank takes the column stripe + ``[tp_rank*(shard_inter//divisor) : +(shard_inter//divisor)]``. Requires + ``shard_inter % divisor == 0`` (asserted at block build) so the stripe lands on + an int32 / group boundary. + """ + assert loaded_shard_id is not None + expert_idx = int(str(loaded_shard_id).split(":")[1]) + shard_inter = divide(full_inter, tp_size) + span = divide(shard_inter, divisor) + start = tp_rank * span + param.data[expert_idx, :, :] = loaded_weight[:, start:start + span] + class KimiMoEGate(nn.Module): """DeepSeek-V3 group-limited sigmoid router with ``noaux_tc`` bias. @@ -170,9 +233,8 @@ class KimiSparseMoeBlock(nn.Module): group, so it shards its intermediate and all-reduces internally. This reuses the existing fused-expert machinery verbatim and is trivially goldenable (tp>1 == tp=1). Its tradeoff: every rank still stores ALL experts' weights, so - it does NOT reduce per-rank expert memory — the 1T fit needs true - token-dispatch expert parallelism (all-to-all), which is a Phase-4 concern and - is deliberately not built here. + it does NOT reduce per-rank expert memory — the 1T fit needs true token-dispatch + expert parallelism (all-to-all), which is deliberately not built here. Expert weights use the fused layout reused from ``model.components.moe``, sharded to this rank (``full == moe_intermediate_size``, @@ -196,6 +258,13 @@ def __init__( # Per-rank slice of each expert's SwiGLU intermediate (== full at tp=1). shard_inter = divide(config.moe_intermediate_size, self.tp_size) + # Packed experts (in-kernel W4A16 dequant) are used iff the checkpoint is + # quantized AND the config opts in. When off, the experts use the bf16 fused + # params (dequantized on load, or a native-bf16 checkpoint loaded directly). + self.packed_experts = ( + config.quantization_config is not None and config.moe_in_kernel_dequant + ) + self.gate = KimiMoEGate( hidden_size=config.hidden_size, n_routed_experts=config.n_routed_experts, @@ -209,27 +278,66 @@ def __init__( ) self.experts = nn.Module() - self.experts.gate_up_proj = nn.Parameter( - torch.empty( - config.n_routed_experts, - 2 * shard_inter, - config.hidden_size, + if self.packed_experts: + # PACKED expert params (int32 weights + bf16 group scales) INSTEAD of the + # bf16 fused params. Layout mirrors the fused bf16 shapes with the K + # (input) axis compressed: gate_up packs K=hidden, down packs K=inter. + # gate_up_proj_packed: int32 (E, 2*shard_inter, hidden // pack_factor) + # gate_up_proj_scale: bf16 (E, 2*shard_inter, hidden // group_size) + # down_proj_packed: int32 (E, hidden, shard_inter // pack_factor) + # down_proj_scale: bf16 (E, hidden, shard_inter // group_size) + qc = config.quantization_config + self.group_size = qc.group_size + self.pack_factor = qc.pack_factor # 8 for INT4 + hidden, gs, pf = config.hidden_size, self.group_size, self.pack_factor + # The packed/group axes must divide evenly on BOTH the hidden (gate_up K) + # and the per-rank intermediate stripe (down K, TP-sharded). + assert hidden % pf == 0 and hidden % gs == 0, ( + f"hidden {hidden} must divide pack_factor {pf} and group_size {gs}" ) - ) - self.experts.down_proj = nn.Parameter( - torch.empty( - config.n_routed_experts, - config.hidden_size, - shard_inter, + assert shard_inter % pf == 0 and shard_inter % gs == 0, ( + f"shard_inter {shard_inter} must divide pack_factor {pf} and group_size {gs}" + ) + E = config.n_routed_experts + self.experts.gate_up_proj_packed = nn.Parameter( + torch.empty(E, 2 * shard_inter, hidden // pf, dtype=torch.int32), + requires_grad=False, + ) + self.experts.gate_up_proj_scale = nn.Parameter( + torch.empty(E, 2 * shard_inter, hidden // gs, dtype=torch.bfloat16), + requires_grad=False, + ) + self.experts.down_proj_packed = nn.Parameter( + torch.empty(E, hidden, shard_inter // pf, dtype=torch.int32), + requires_grad=False, + ) + self.experts.down_proj_scale = nn.Parameter( + torch.empty(E, hidden, shard_inter // gs, dtype=torch.bfloat16), + requires_grad=False, + ) + else: + self.experts.gate_up_proj = nn.Parameter( + torch.empty( + config.n_routed_experts, + 2 * shard_inter, + config.hidden_size, + ) + ) + self.experts.down_proj = nn.Parameter( + torch.empty( + config.n_routed_experts, + config.hidden_size, + shard_inter, + ) ) - ) # The fused expert params are plain nn.Parameters, so they carry no - # per-shard ``weight_loader`` by default. The M5 stacked-param rules route - # each checkpoint expert via a ``"gate:N"/"up:N"/"down:N"`` shard id, so we - # attach the same fused-expert loaders ``ParallelSparseMoeBlock`` uses. - # The loaders take ``(tp_rank, tp_size, full_inter)`` and slice this rank's - # intermediate stripe out of the full-size checkpoint expert — so a single - # weight path serves tp=1 (full) and tp>1 (sharded). + # per-shard ``weight_loader`` by default. The stacked-param rules route each + # checkpoint expert via a ``"gate:N"/"up:N"/"down:N"`` shard id, so we attach + # the same fused-expert loaders ``ParallelSparseMoeBlock`` uses (or, when + # packed, the packed analogs). The loaders take + # ``(tp_rank, tp_size, full_inter)`` and slice this rank's intermediate + # stripe out of the full-size checkpoint expert — so a single weight path + # serves tp=1 (full) and tp>1 (sharded). self._attach_expert_weight_loaders() # Ungated shared expert: a dense SwiGLU MLP with the shared intermediate @@ -248,16 +356,34 @@ def _attach_expert_weight_loaders(self) -> None: Mirrors ``ParallelSparseMoeBlock._attach_weight_loaders``. Re-run after every ``_apply`` (``.to(dtype)`` / ``to_empty(device)`` rebuild the Parameter objects and drop the attribute), so weights load correctly - through the meta -> to_empty -> load path. + through the meta -> to_empty -> load path. When packed, the four packed / + scale params get the Kimi-local packed loaders; otherwise the two bf16 + fused params get the shared loaders. """ from functools import partial - self.experts.gate_up_proj.weight_loader = partial( - _gate_up_weight_loader, self.tp_rank, self.tp_size, self.moe_intermediate_size, - ) - self.experts.down_proj.weight_loader = partial( - _down_proj_weight_loader, self.tp_rank, self.tp_size, self.moe_intermediate_size, - ) + full_inter = self.moe_intermediate_size + if self.packed_experts: + pf, gs = self.pack_factor, self.group_size + self.experts.gate_up_proj_packed.weight_loader = partial( + _gate_up_packed_loader, self.tp_rank, self.tp_size, full_inter, + ) + self.experts.gate_up_proj_scale.weight_loader = partial( + _gate_up_packed_loader, self.tp_rank, self.tp_size, full_inter, + ) + self.experts.down_proj_packed.weight_loader = partial( + _down_packed_loader, self.tp_rank, self.tp_size, full_inter, pf, + ) + self.experts.down_proj_scale.weight_loader = partial( + _down_packed_loader, self.tp_rank, self.tp_size, full_inter, gs, + ) + else: + self.experts.gate_up_proj.weight_loader = partial( + _gate_up_weight_loader, self.tp_rank, self.tp_size, full_inter, + ) + self.experts.down_proj.weight_loader = partial( + _down_proj_weight_loader, self.tp_rank, self.tp_size, full_inter, + ) def _apply(self, fn, recurse=True): result = super()._apply(fn, recurse=recurse) @@ -271,7 +397,11 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # Router is replicated: every rank computes the full top-k selection. topk_weights, topk_ids = self.gate(flat) topk_weights = topk_weights.to(flat.dtype) - if self.tp_size == 1: + if self.packed_experts: + # Packed experts: bypass the shared bf16 ``_dispatch`` and run the + # W4A16 in-kernel dequant GEMM directly (handles tp=1 and tp>1). + routed = self._dispatch_packed_experts(flat, topk_weights, topk_ids) + elif self.tp_size == 1: routed = _dispatch( flat, self.experts.gate_up_proj, @@ -287,6 +417,45 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: shared = self.shared_expert(flat) return (routed + shared).view(input_shape) + def _dispatch_packed_experts( + self, + flat: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + ) -> torch.Tensor: + """Packed (W4A16) routed dispatch — the memory-lean packed-expert path. + + Feeds the packed int32 weights + bf16 group scales to ``fused_experts``, + which launches ``fused_moe_kernel_w4a16`` (dequant in registers). The TP + story is identical to :meth:`_dispatch_tp`: at tp=1 the kernel sum-reduces + the top-k dim itself; at tp>1 we keep the per-slot partials + (``reduce_results=False``), all-reduce the intermediate-parallel partials, + then fold the top-k dim. Combine weights already carry + ``routed_scaling_factor`` (folded by :class:`KimiMoEGate`), so the + sum-reduce passes ``routed_scaling_factor=1.0``. + """ + from mstar.utils.fused_moe import fused_experts, moe_sum_reduce_triton + + reduce = self.tp_size == 1 + out = fused_experts( + flat, + self.experts.gate_up_proj_packed, + self.experts.down_proj_packed, + topk_weights, + topk_ids, + w1_scale=self.experts.gate_up_proj_scale, + w2_scale=self.experts.down_proj_scale, + group_size=self.group_size, + pack_factor=self.pack_factor, + reduce_results=reduce, + ) + if reduce: + return out + self.comm_group.all_reduce(out) + output = torch.empty_like(flat) + moe_sum_reduce_triton(out, output, routed_scaling_factor=1.0) + return output + def _dispatch_tp( self, flat: torch.Tensor, diff --git a/mstar/model/kimi_k2_7/components/rope.py b/mstar/model/kimi_k2_7/components/rope.py index 3cde1b4f9..e170af20f 100644 --- a/mstar/model/kimi_k2_7/components/rope.py +++ b/mstar/model/kimi_k2_7/components/rope.py @@ -3,13 +3,11 @@ MLA rotates only the decoupled ``qk_rope_head_dim`` slice of q/k, with YARN (NTK-by-parts) frequency scaling and an ``mscale`` amplitude on cos/sin. mstar's ``cache_manager.apply_rope`` (FlashInfer) does not implement YARN, so this is a -standalone rotary module the MLA attention applies itself (analogous to how -Qwen3-Omni applies its 3D MRoPE outside the cache handle). +standalone rotary module the MLA attention applies itself. Style is **interleaved / GPT-J** (``is_neox_style=False`` in DeepSeek): cos/sin are ``repeat_interleave(2)`` and adjacent even/odd pairs are rotated. Mirrors -vLLM ``layers/rotary_embedding/deepseek_scaling_rope.py::DeepseekScalingRotaryEmbedding`` -and the YARN helpers in ``rotary_embedding/common.py``. +vLLM ``DeepseekScalingRotaryEmbedding`` and the YARN helpers in ``common.py``. Two ``mscale`` values (both use the 2-arg ``yarn_get_mscale``): - **amplitude** on cos/sin (here): ``get_mscale(f, mscale) / get_mscale(f, mscale_all_dim) * attn_factor``. diff --git a/mstar/model/kimi_k2_7/config.py b/mstar/model/kimi_k2_7/config.py index 798f0f9a5..e4041c38f 100644 --- a/mstar/model/kimi_k2_7/config.py +++ b/mstar/model/kimi_k2_7/config.py @@ -1,25 +1,22 @@ """Configuration dataclass for Kimi-K2.7 (text backbone). Kimi-K2.7's text architecture *is* DeepSeek-V3 — vLLM serves it as -``DeepseekV3ForCausalLM`` (``model_type: "kimi_k2"`` maps to -``DeepseekV3Config``). This dataclass therefore carries the full DeepSeek-V3 -field set: MLA latent dims, fine-grained sigmoid-routed MoE grouping, and -``deepseek_yarn`` RoPE. Only a handful of these fields are read by the M0 -scaffold (``num_hidden_layers``, the head dims, ``vocab_size``, -``max_position_embeddings``); the rest are declared now so this stays the single -source of truth for the later milestones (MoE router, MLA attention, weights). - -The full-size defaults below are the real values from the -``moonshotai/Kimi-K2.7-Code`` HF ``config.json`` (``model_type: "kimi_k2"`` → -``DeepseekV3Config``), confirmed field by field against the published checkpoint -config (the weights themselves are not present in this workspace). M0 does not -depend on the full-size values being exact — the modular tests build from -:meth:`KimiK2Config.reduced`, a tiny self-consistent config. +``DeepseekV3ForCausalLM`` (``model_type: "kimi_k2"`` -> ``DeepseekV3Config``), so +this dataclass carries the full DeepSeek-V3 field set: MLA latent dims, sigmoid- +routed MoE grouping, and ``deepseek_yarn`` RoPE. + +The full-size defaults are the real ``moonshotai/Kimi-K2.7-Code`` values. That +repo is the multimodal ``KimiK25ForConditionalGeneration``; the text dims here +live NESTED under ``config.json``'s ``text_config``, and its ``quantization_config`` +is nested there too (see :meth:`k27_code` and ``kimi_model.py``). The modular tests +build from :meth:`reduced`, a tiny self-consistent config. """ from __future__ import annotations from dataclasses import dataclass, field +from mstar.model.kimi_k2_7.quantization import CompressedTensorsQuantConfig + @dataclass class KimiK2Config: @@ -65,8 +62,8 @@ class KimiK2Config: rope_scaling: dict = field(default_factory=lambda: { # from config.json (HF key is "type": "yarn"; mstar's internal id for the # DeepSeek/Kimi variant is "deepseek_yarn"). factor=64 yields the 262144 - # context (4096 * 64). K2.7-Code keeps beta_fast=32 (some other Kimi - # checkpoints set beta_fast=1). mscale == mscale_all_dim == 1.0. + # context (4096 * 64). K2.7-Code keeps beta_fast=32; mscale == + # mscale_all_dim == 1.0. "rope_type": "deepseek_yarn", "factor": 64.0, "original_max_position_embeddings": 4096, @@ -87,11 +84,28 @@ class KimiK2Config: # -- MTP (multi-token prediction) — deferred, declared for completeness - num_nextn_predict_layers: int = 0 + # -- Quantization (compressed-tensors INT4/fp8) ----------------------- + # ``None`` => native-bf16 checkpoint. When set, the weight loader dequantizes + # the checkpoint stream on load (:mod:`mstar.model.kimi_k2_7.quantization`) + # before the bf16 remap + stacked rules. Populated from the real checkpoint's + # ``config.json`` ``quantization_config`` (``kimi_model.py``) or set directly + # for the reduced/synthetic tests (:meth:`reduced_quantized`). + quantization_config: CompressedTensorsQuantConfig | None = None + + # -- Quantization: memory-lean packed experts (in-kernel dequant) ---------- + # ``False`` => quantized routed experts are dequantized to bf16 on load and fed + # to the bf16 fused-expert GEMM. ``True`` (only meaningful when + # ``quantization_config`` is set) => the routed experts stay PACKED int32 in + # VRAM and the W4A16 ``fused_moe_kernel_w4a16`` dequantizes each tile in + # registers. MLA / dense-FFN / shared-expert weights are always dequantized on + # load. This is the only path that fits the real 1T checkpoint. See + # ``components/moe.py`` / ``weight_loader.py``. + moe_in_kernel_dequant: bool = False + # -- Serving: CUDA-graph prefill capture grid (optional overrides) ------ # ``None`` => ``KimiLLMSubmodule`` uses its full-size class-default grid. # ``reduced()`` sets a tiny grid so the synthetic bring-up serve captures a - # single short-prompt graph instead of the full 6x5 compiled grid (this was - # an env knob during bring-up; it now lives in the config). + # single short-prompt graph instead of the full 6x5 compiled grid. prefill_token_buckets: list[int] | None = None prefill_capture_batch_sizes: list[int] | None = None @@ -108,14 +122,12 @@ def padded_head_dim(self) -> int: """Head dim the naive-MLA q/k/v are zero-padded to for the paged cache. FlashInfer's SM90 (Hopper) prefill kernel ``static_assert``s - ``head_dim_vo ∈ {64, 128, 256}`` (M4 finding), so it will not JIT-build for - the real ``qk_head_dim=192`` or the reduced ``qk_head_dim=24``. The - correctness-first mitigation (M6) pads q/k (from ``qk_head_dim``) and v - (from ``v_head_dim``) up to the smallest supported dim ``>= qk_head_dim``, - runs the paged attention there, and slices the output back to - ``v_head_dim`` — compensating the softmax scale (see - ``KimiMLAAttention.softmax_scale_boost``). Real Kimi 192 -> 256; reduced - 24 -> 64. Weight-absorbed MLA (which avoids the pad) is deferred to perf. + ``head_dim_vo ∈ {64, 128, 256}``, so it will not JIT-build for the real + ``qk_head_dim=192`` or the reduced ``qk_head_dim=24``. We pad q/k (from + ``qk_head_dim``) and v (from ``v_head_dim``) up to the smallest supported + dim ``>= qk_head_dim``, run the paged attention there, and slice the output + back to ``v_head_dim`` — compensating the softmax scale (see + ``KimiMLAAttention.softmax_scale_boost``). Real Kimi 192 -> 256; reduced 24 -> 64. """ for supported in (64, 128, 256): if supported >= self.qk_head_dim: @@ -163,3 +175,66 @@ def reduced(cls) -> "KimiK2Config": prefill_token_buckets=[64], prefill_capture_batch_sizes=[1], ) + + @classmethod + def reduced_quantized( + cls, + num_bits: int = 4, + group_size: int = 32, + symmetric: bool = True, + ) -> "KimiK2Config": + """:meth:`reduced` plus a compressed-tensors quant config, to exercise the + dequant-on-load path on a synthetic quantized checkpoint. + + The reduced dims (``hidden_size=128``, ``moe_intermediate_size=64``, + ``intermediate_size=256`` …) are all divisible by the default + ``group_size=32`` and by ``pack_factor=8``, so the FFN / expert / MLA + weights whose input dim divides ``group_size`` can be quantized while the + rest stay bf16 — the mixed checkpoint the streaming parser handles. + """ + cfg = cls.reduced() + cfg.quantization_config = CompressedTensorsQuantConfig( + num_bits=num_bits, group_size=group_size, symmetric=symmetric, + ) + return cfg + + @classmethod + def reduced_quantized_inkernel( + cls, + num_bits: int = 4, + group_size: int = 32, + symmetric: bool = True, + ) -> "KimiK2Config": + """:meth:`reduced_quantized` plus ``moe_in_kernel_dequant=True`` — packed + routed experts + in-kernel INT4 dequant on a synthetic quantized checkpoint. + + The reduced dims (``hidden_size=128``, ``moe_intermediate_size=64``) satisfy + the packed-expert divisibility asserts (``% pack_factor`` and ``% + group_size``) at tp=1 (``shard_inter=64``) and tp=2 (``shard_inter=32``); + tp=4 (``shard_inter=16``) fails ``% group_size`` (32), so pin packed-expert + goldens to tp<=2. + """ + cfg = cls.reduced_quantized( + num_bits=num_bits, group_size=group_size, symmetric=symmetric, + ) + cfg.moe_in_kernel_dequant = True + return cfg + + @classmethod + def k27_code(cls) -> "KimiK2Config": + """Full-size ``moonshotai/Kimi-K2.7-Code`` text-only serve config. + + Full 1T dims plus ``moe_in_kernel_dequant=True``: Kimi-K2.7-Code is a ~1T + INT4 ``pack-quantized`` checkpoint (num_bits=4, group_size=32, symmetric, + routed experts only), so the routed experts are served packed and + dequantized in-kernel — dequantizing them to bf16 would need ~2 TB of VRAM. + The ``quantization_config`` is nested under ``text_config`` and auto-read at + load by ``kimi_model.py::_maybe_apply_checkpoint_quant_config``; MLA / + dense-FFN / shared-expert / lm_head / vision weights stay bf16, matching the + checkpoint ``ignore`` list. + + Keeps the default ``beta_fast=32.0``. + """ + cfg = cls() + cfg.moe_in_kernel_dequant = True + return cfg diff --git a/mstar/model/kimi_k2_7/kimi_model.py b/mstar/model/kimi_k2_7/kimi_model.py index 9978d300c..3e0b9ae77 100644 --- a/mstar/model/kimi_k2_7/kimi_model.py +++ b/mstar/model/kimi_k2_7/kimi_model.py @@ -1,22 +1,17 @@ """KimiK2Model: M* Model contract for Kimi-K2.7 (text backbone). -Kimi-K2.7's text path is DeepSeek-V3 (``model_type: "kimi_k2"`` → -``DeepseekV3ForCausalLM``). This is the **M0 scaffold**: it declares the full -serving plumbing — the graph (``prefill`` + ``decode`` Loop), the single -``KV_CACHE`` LLM node, the KV-cache dims, and the prefill→decode→done state -machine — with **no GPU compute**. ``get_submodule`` returns ``None`` (dummy -mode), so ``pytest test/modular/`` exercises the graph/walk/engine-routing -machinery in isolation, exactly as ``docs/adding_models.rst`` prescribes for a -new model before touching weights. +Kimi-K2.7's text path is DeepSeek-V3 (``model_type: "kimi_k2"`` -> +``DeepseekV3ForCausalLM``). This declares the full serving plumbing — the graph +(``prefill`` + ``decode`` Loop), the single ``KV_CACHE`` LLM node, the KV-cache +dims, and the prefill->decode->done state machine — and builds the LLM submodule +in ``get_submodule``. When ``get_submodule`` returns ``None`` (dummy mode), +``pytest test/modular/`` exercises the graph/walk/engine-routing machinery in +isolation, as ``docs/adding_models.rst`` prescribes. Structurally this mirrors Orpheus's LLM partition (the smallest complete LLM in the tree) minus the async SNAC partition: Kimi text-only is a single ``default`` partition, so it inherits ``Model.get_partitions`` / ``get_partition_topology`` and only implements the abstract surface. - -Later milestones fill in the real compute (M2 MoE router, M3 MLA attention, M5 -weights, M6 the ``KimiLLMSubmodule`` build in ``get_submodule``); none of them -change the contract declared here. """ from __future__ import annotations @@ -87,6 +82,17 @@ def __init__( self._config_variant = kwargs.get("config_variant", "full") if self._config_variant == "reduced": self.config = KimiK2Config.reduced() + elif self._config_variant == "reduced_quantized": + # Reduced shape + a quant config, to exercise dequant-on-load. + self.config = KimiK2Config.reduced_quantized() + elif self._config_variant == "reduced_quantized_inkernel": + # Reduced shape + quant config + packed experts (in-kernel W4A16 dequant). + # int32 packed params are auto-exempt from the whole-model ``.to(bf16)`` + # cast below (PyTorch ``.to(dtype)`` only casts float/complex), no hook. + self.config = KimiK2Config.reduced_quantized_inkernel() + elif self._config_variant == "k27_code": + # Full-size Kimi-K2.7-Code text-only serve config (see KimiK2Config.k27_code). + self.config = KimiK2Config.k27_code() else: self.config = KimiK2Config() self._tokenizer_mode = kwargs.get("tokenizer_mode", "hf") @@ -120,10 +126,9 @@ def get_kv_cache_config(self) -> list[KVCacheConfig]: # ``padded_head_dim`` — the naive path zero-pads q/k (from ``qk_head_dim``, # e.g. 192) and v (from ``v_head_dim``) up to the smallest FlashInfer-SM90 # supported head_dim >= qk_head_dim (256 real, 64 reduced), because the - # Hopper prefill kernel static_asserts head_dim_vo in {64,128,256} (M4 - # finding). The attention output is sliced back to ``v_head_dim`` in the - # submodule (M3/M6). This trades cache size for not needing a weight-absorb - # path in the engine (deferred to perf). + # Hopper prefill kernel static_asserts head_dim_vo in {64,128,256}. The + # attention output is sliced back to ``v_head_dim`` in the submodule. This + # trades cache size for not needing a weight-absorb path in the engine. return [KVCacheConfig( num_layers=self.config.num_hidden_layers, num_kv_heads=self.config.num_attention_heads, @@ -337,7 +342,7 @@ def get_default_sharding_config(self): # Kimi is a 1T model — real serving is TP8 / multi-node. The LLM node is # the tensor-parallel node; the per-node degree comes from the config - # YAML's ``node_groups`` (M6), not from the model code. + # YAML's ``node_groups``, not from the model code. return ShardingConfig(groups=[], tp_enabled_nodes={LLM_NODE}, shard_dim={}) # ------------------------------------------------------------------- @@ -382,12 +387,15 @@ def _create_submodule( ) return None + # If the checkpoint declares a compressed-tensors ``quantization_config``, + # route loading through the dequant-on-load parser (weight_loader). An + # explicit config (e.g. reduced_quantized) is respected and not clobbered. + self._maybe_apply_checkpoint_quant_config(source) + # Real build, mirroring OrpheusModel._create_llm_submodule: construct on the # meta device (no allocation), cast to the target dtype on meta (so to_empty # allocates directly in bf16, not fp32-then-downcast), materialise storage, - # then run the M5 HF loader (remap + fused-expert stacked rules). This is the - # path the M5 rope-buffer bug would have bitten — inv_freq is lazy so it does - # not survive as garbage. + # then run the HF loader (remap + fused-expert stacked rules). from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM from mstar.model.kimi_k2_7.submodules import KimiLLMSubmodule from mstar.model.loader import load_weights @@ -418,3 +426,47 @@ def _resolve_checkpoint(self) -> str | None: if Path(path).exists(): return str(path) return _resolve_local_hf_snapshot(path, cache_dir=getattr(self, "cache_dir", None)) + + def _maybe_apply_checkpoint_quant_config(self, source: str) -> None: + """Populate ``self.config.quantization_config`` from ``config.json``. + + Reads the checkpoint's ``config.json`` ``quantization_config`` block (a + compressed-tensors INT4/fp8 checkpoint carries one) and stores the parsed + :class:`CompressedTensorsQuantConfig` on the model config so the weight + loader takes the dequant-on-load path. A config set explicitly + (e.g. ``reduced_quantized()``) wins and is left untouched; a plain bf16 + checkpoint (no block, unreadable, or single-file source) is a no-op. + + The real multimodal ``Kimi-K2.7-Code`` repo nests the block under + ``text_config`` (the top-level ``quantization_config`` is null), so this + reads a top-level block if present, otherwise ``text_config``'s. + ``from_hf_config_dict`` parses the same block shape either way. + """ + import json + from pathlib import Path + + from mstar.model.kimi_k2_7.quantization import CompressedTensorsQuantConfig + + if self.config.quantization_config is not None: + return + config_json = Path(source) / "config.json" + if not config_json.is_file(): + return + try: + with open(config_json) as f: + raw = json.load(f) + except (OSError, ValueError) as e: # unreadable / malformed — stay bf16 + logger.warning("KimiK2Model: could not read %s: %s", config_json, e) + return + # A top-level ``quantization_config`` if present, else the one nested under + # ``text_config`` (the multimodal K2.7-Code layout — top-level is null). + quant_raw = raw.get("quantization_config") or ( + raw.get("text_config") or {} + ).get("quantization_config") + quant = CompressedTensorsQuantConfig.from_hf_config_dict(quant_raw) + if quant is not None: + logger.info( + "KimiK2Model: compressed-tensors checkpoint (%d-bit, group_size=%d) " + "— dequantizing on load.", quant.num_bits, quant.group_size, + ) + self.config.quantization_config = quant diff --git a/mstar/model/kimi_k2_7/quantization.py b/mstar/model/kimi_k2_7/quantization.py new file mode 100644 index 000000000..d0ff5085e --- /dev/null +++ b/mstar/model/kimi_k2_7/quantization.py @@ -0,0 +1,245 @@ +"""Compressed-tensors INT4 (W4A16) quantization for Kimi-K2.7 weights. + +On-disk format (compressed-tensors ``pack-quantized``), per quantized Linear +weight of logical shape ``(out, in)``: + + * ``.weight_packed`` — int32 ``(out, in // pack_factor)``; ``pack_factor = + 32 // num_bits`` (8 for INT4) values packed low-order-first along the input axis. + * ``.weight_scale`` — bf16 ``(out, in // group_size)``, one scale per (row, group). + * ``.weight_zero_point`` — asymmetric only. + * ``.weight_shape`` — original ``(out, in)``; optional, used only to validate. + +Symmetric INT4 is stored offset-binary: the packed nibble is ``(signed value + 8)``, +so dequant subtracts 8 (matches vLLM's ``uint4b8``); asymmetric subtracts the zero +point. To flip a checkpoint to plain two's-complement, change the ``bias`` line in +:func:`dequantize_weight`. Layout authority: vLLM +``compressed_tensors/schemes/compressed_tensors_wNa16.py``. + +``dequant_compressed_tensors_stream`` dequantizes a checkpoint stream to bf16 on +load; ``keep_packed`` leaves selected weights packed for in-kernel dequant. +""" +from __future__ import annotations + +from collections.abc import Callable, Iterable, Iterator +from dataclasses import dataclass, field + +import torch + +# Suffixes a compressed-tensors checkpoint attaches to each quantized tensor. +_PACKED = ".weight_packed" +_SCALE = ".weight_scale" +_ZERO_POINT = ".weight_zero_point" +_SHAPE = ".weight_shape" +_QUANT_SUFFIXES = (_PACKED, _SCALE, _ZERO_POINT, _SHAPE) + + +@dataclass(frozen=True) +class CompressedTensorsQuantConfig: + """The subset of a compressed-tensors ``quantization_config`` this port reads. + + Kimi-K2.7 ships a single ``config_groups`` entry (``weights`` only, W4A16), so + the whole checkpoint shares one ``num_bits`` / ``group_size`` / ``symmetric``. + """ + + num_bits: int = 4 + group_size: int = 32 # -1 => channelwise (one group spans the full input dim) + symmetric: bool = True + strategy: str = "group" # "group" | "channel" + quant_format: str = "pack-quantized" + quant_method: str = "compressed-tensors" + ignore: tuple[str, ...] = field(default_factory=tuple) + + @property + def pack_factor(self) -> int: + """Number of ``num_bits`` values packed into one int32.""" + return 32 // self.num_bits + + @classmethod + def from_hf_config_dict( + cls, quant: dict | None + ) -> "CompressedTensorsQuantConfig | None": + """Build from a checkpoint ``config.json``'s ``quantization_config`` block. + + Returns ``None`` when there is no quantization block (a plain bf16 + checkpoint). Reads the first ``config_groups`` entry's ``weights`` spec — + Kimi uses exactly one group. + """ + if not quant: + return None + groups = quant.get("config_groups") or {} + weights: dict = {} + if groups: + first = next(iter(groups.values())) + weights = (first or {}).get("weights") or {} + strategy = weights.get("strategy", "group") + group_size = weights.get("group_size") + if group_size is None: + group_size = -1 if strategy == "channel" else 32 + return cls( + num_bits=int(weights.get("num_bits", 4)), + group_size=int(group_size), + symmetric=bool(weights.get("symmetric", True)), + strategy=str(strategy), + quant_format=str(quant.get("format", "pack-quantized")), + quant_method=str(quant.get("quant_method", "compressed-tensors")), + ignore=tuple(quant.get("ignore", []) or []), + ) + + +# --------------------------------------------------------------------------- +# Bit packing — exact inverses. Packing is along the last (input) axis, which is +# the input axis of a checkpoint ``(out, in)`` Linear weight (what gets quantized). +# --------------------------------------------------------------------------- + +def pack_int32(values_unsigned: torch.Tensor, num_bits: int) -> torch.Tensor: + """Pack ``pack_factor`` unsigned ``num_bits`` values (last axis) into int32. + + ``values_unsigned`` holds integers in ``[0, 2**num_bits)``; the last axis must + be divisible by ``pack_factor = 32 // num_bits``. Values are combined + low-order-first (element ``j`` occupies bits ``[num_bits*j, num_bits*(j+1))``), + matching compressed-tensors. Returns int32 of shape + ``(..., last // pack_factor)``. + """ + pack_factor = 32 // num_bits + *lead, n = values_unsigned.shape + if n % pack_factor != 0: + raise ValueError(f"last dim {n} not divisible by pack_factor {pack_factor}") + q = values_unsigned.to(torch.int64).reshape(*lead, n // pack_factor, pack_factor) + shifts = torch.arange(pack_factor, device=q.device, dtype=torch.int64) * num_bits + packed = (q << shifts).sum(dim=-1) + # Wrap the 32-bit pattern into a signed int32 container (matches on-disk dtype). + return (packed & 0xFFFFFFFF).to(torch.int32) + + +def unpack_int32(packed: torch.Tensor, num_bits: int) -> torch.Tensor: + """Inverse of :func:`pack_int32`: expand int32 to unsigned ``num_bits`` nibbles. + + Returns an int64 tensor of shape ``(..., last * pack_factor)`` with values in + ``[0, 2**num_bits)``. The int32 is read as an unsigned 32-bit pattern, so the + top nibble is recovered correctly regardless of the container's sign bit. + """ + pack_factor = 32 // num_bits + mask = (1 << num_bits) - 1 + p = packed.to(torch.int64) & 0xFFFFFFFF + shifts = torch.arange(pack_factor, device=p.device, dtype=torch.int64) * num_bits + unpacked = (p.unsqueeze(-1) >> shifts) & mask # (..., last, pack_factor) + *lead, m, _ = unpacked.shape + return unpacked.reshape(*lead, m * pack_factor) + + +def dequantize_weight( + packed: torch.Tensor, + scale: torch.Tensor, + *, + num_bits: int, + group_size: int, + symmetric: bool, + zero_point: torch.Tensor | None = None, + out_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Dequantize one compressed-tensors weight to ``out_dtype``. + + Args: + packed: ``(out, in // pack_factor)`` int32 packed weight. + scale: ``(out, in // group_size)`` per-(row, group) scale. + num_bits: bit width (4 for Kimi INT4). + group_size: group granularity along the input axis; ``-1`` => channelwise. + symmetric: symmetric offset-binary (subtract ``2**(num_bits-1)``) vs. + asymmetric (subtract ``zero_point``). + zero_point: ``(out, in // group_size)`` per-group zero point (asymmetric). + out_dtype: result dtype (bf16 to feed the existing fused-expert GEMM). + + Returns: + ``(out, in)`` dequantized weight, in ``out_dtype``. + """ + nibbles = unpack_int32(packed, num_bits).to(torch.float32) # (out, in) unsigned + out_f, in_f = nibbles.shape + gs = in_f if group_size in (-1, None) else group_size + if in_f % gs != 0: + raise ValueError(f"in_features {in_f} not divisible by group_size {gs}") + + if symmetric: + nibbles -= float(1 << (num_bits - 1)) # offset-binary: nibble - bias + else: + if zero_point is None: + raise ValueError("asymmetric quantization requires a zero_point") + zp = zero_point.to(torch.float32) + if zp.shape[-1] != in_f: # per-group -> broadcast to per-column + zp = zp.repeat_interleave(gs, dim=-1) + nibbles -= zp + + s = scale.to(torch.float32) + if s.shape[-1] != in_f: # per-group -> broadcast to per-column + s = s.repeat_interleave(gs, dim=-1) + return (nibbles * s).to(out_dtype) + +# --------------------------------------------------------------------------- +# Streaming dequant-on-load — the generator wired into load_kimi_hf_weights. +# --------------------------------------------------------------------------- + +def dequant_compressed_tensors_stream( + weights: Iterable[tuple[str, torch.Tensor]], + quant_config: CompressedTensorsQuantConfig, + out_dtype: torch.dtype = torch.bfloat16, + keep_packed: Callable[[str], bool] | None = None, +) -> Iterator[tuple[str, torch.Tensor]]: + """Wrap a checkpoint ``(name, tensor)`` stream, dequantizing on the fly. + + For every quantized tensor the checkpoint carries ``.weight_packed`` + + ``.weight_scale`` (+ ``.weight_zero_point`` for asymmetric); this + buffers those components per ```` and, once complete, yields a single + ``(.weight, bf16 tensor)`` — exactly the key a native-bf16 checkpoint + would carry — then drops the quant sub-keys. Any key that is not a + compressed-tensors component (norms, the router ``gate``, ``embed_tokens``, + ``lm_head``, or a weight the checkpoint left in bf16) passes straight through. + + Buffering is bounded to the in-flight incomplete tensors: a ```` is + emitted and freed the moment its required components have all been seen, + independent of the iterator's key order. + + ``keep_packed``: when it returns True for a ````, that base's + compressed-tensors sub-keys are passed through RAW (no buffering, no dequant) + so a downstream packed-expert loader can route them to int32 params. Kimi + passes a predicate matching the routed experts (kept packed for in-kernel + dequant) while every other quantized weight — MLA, dense FFN, shared expert — + still dequantizes here. Because kept bases never enter ``buffers``, the + end-of-stream completeness check is unaffected. + """ + buffers: dict[str, dict[str, torch.Tensor]] = {} + + for name, tensor in weights: + suffix = next((s for s in _QUANT_SUFFIXES if name.endswith(s)), None) + if suffix is None: + yield name, tensor # not a quant component — pass through untouched + continue + + base = name[: -len(suffix)] + if keep_packed is not None and keep_packed(base): + # Leave this base packed — hand the raw sub-key downstream. + yield name, tensor + continue + + slot = buffers.setdefault(base, {}) + slot[suffix] = tensor + + have_core = _PACKED in slot and _SCALE in slot + have_zp = quant_config.symmetric or _ZERO_POINT in slot + if have_core and have_zp: + weight = dequantize_weight( + slot[_PACKED], + slot[_SCALE], + num_bits=quant_config.num_bits, + group_size=quant_config.group_size, + symmetric=quant_config.symmetric, + zero_point=slot.get(_ZERO_POINT), + out_dtype=out_dtype, + ) + del buffers[base] + yield base + ".weight", weight + + if buffers: + missing = {b: sorted(slot) for b, slot in buffers.items()} + raise ValueError( + f"compressed-tensors stream ended with incomplete quantized tensors " + f"(missing weight_packed and/or weight_scale): {missing}" + ) diff --git a/mstar/model/kimi_k2_7/submodules.py b/mstar/model/kimi_k2_7/submodules.py index ed6bef6c2..49294b042 100644 --- a/mstar/model/kimi_k2_7/submodules.py +++ b/mstar/model/kimi_k2_7/submodules.py @@ -1,8 +1,8 @@ """Submodules for Kimi-K2.7 (text backbone). -M6: the real :class:`KimiLLMSubmodule` — the ``ARNodeSubmodule`` that drives the -DeepSeek-V3 text backbone (MLA attention over the paged cache + fine-grained -sigmoid-routed MoE) through the engine's ``prepare_inputs -> preprocess -> +The :class:`KimiLLMSubmodule` ``ARNodeSubmodule`` drives the DeepSeek-V3 text +backbone (MLA attention over the paged cache + fine-grained sigmoid-routed MoE) +through the engine's ``prepare_inputs -> preprocess -> forward/forward_batched -> postprocess -> check_stop`` lifecycle for the ``prefill`` and ``decode`` Loop walks. @@ -59,10 +59,10 @@ def __init__(self, language_model: nn.Module, config: KimiK2Config): # -- CUDA-graph prefill capture grid (full-size defaults) -------------- # Overridable per-config via ``config.prefill_token_buckets`` / # ``config.prefill_capture_batch_sizes``: ``KimiK2Config.reduced()`` sets a - # tiny grid for the synthetic bring-up serve, while the full model leaves them + # tiny grid for a reduced-size serve, while the full model leaves them # ``None`` and uses these defaults. Capturing the full 6x5 compiled grid is # slow, and buckets above a small model's ``max_position_embeddings`` do not - # fit — hence the reduced config trims it (config-driven, not env-driven). + # fit — hence the reduced config trims it. PREFILL_TOKEN_BUCKETS = [32, 64, 128, 256, 512, 1024] PREFILL_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16] diff --git a/mstar/model/kimi_k2_7/weight_loader.py b/mstar/model/kimi_k2_7/weight_loader.py index d80d40537..a7f38e3ab 100644 --- a/mstar/model/kimi_k2_7/weight_loader.py +++ b/mstar/model/kimi_k2_7/weight_loader.py @@ -1,151 +1,156 @@ -"""Kimi-K2.7 / DeepSeek-V3 weight loading (M5). - -Maps an HF ``DeepseekV3ForCausalLM`` checkpoint onto the mstar Kimi module tree -using the shared ``load_hf_weights`` machinery — a name remap plus stacked-shard -rules, exactly mirroring ``qwen3_omni_model.py::_get_thinker_stacked_params`` / -``_thinker_remap`` (same fused-expert layout). No shared abstraction is modified. - -Two transforms take the checkpoint keys to the module's ``named_parameters``: - -1. **Name remap** (:func:`kimi_name_remapper`): - - ``mlp.shared_experts.*`` -> ``mlp.shared_expert.*`` (HF plural -> our - singular ``ParallelGatedMLP`` submodule name); - - per-routed-expert ``mlp.experts.{i}.{gate,up,down}_proj.weight`` -> - ``mlp.experts.{gate,up,down}_proj.__expert{i}__.weight`` — the - ``__expert{i}__`` marker lets the stacked rules carry both the projection - *and* the expert slot in one ``shard_id``; - - everything else is identity (the checkpoint prefixes ``model.``, - ``self_attn.``, ``input_layernorm``, ``lm_head`` … already line up). - -2. **Stacked rules** (:func:`build_kimi_stacked_params`): - - routed experts: ``.experts.gate_proj.__expert{i}__.weight`` / - ``.up_proj.__expert{i}__.weight`` -> fused ``.experts.gate_up_proj`` - ("w13", gate then up) with ``shard_id="gate:i"/"up:i"``; - ``.down_proj.__expert{i}__.weight`` -> ``.experts.down_proj`` ("w2", - ``shard_id="down:i"``). The fused params get their per-shard - ``weight_loader`` in :class:`KimiSparseMoeBlock`. - - dense + shared SwiGLU ``.gate_proj`` / ``.up_proj`` -> merged - ``.gate_up_proj`` (shard 0 / 1). These MUST come *after* the expert rules: - ``_apply_stacked`` returns on first match and the remapped expert key - ``…experts.gate_proj.__expert{i}__.weight`` also contains ``.gate_proj``. - -**MLA loads strictly by name — NO q_a/kv_a fusion.** M3 built *separate* -``q_a_proj`` and ``kv_a_proj_with_mqa`` (the naive/materialized path), so their -checkpoint keys map straight to the identically-named params. Fusing them into a -single ``fused_qkv_a_proj`` is only needed for the deferred weight-absorbed MLA -class (``DeepseekV2MLAAttention``); the naive path needs no such fusion. This is -a deliberate simplification of the earlier plan note. - -**Router bias stays fp32.** ``KimiMoEGate.e_score_correction_bias`` is a -selection-only bias DeepSeek keeps in fp32 for router stability. A whole-model -``.to(bfloat16)`` would downcast it, so :func:`restore_router_bias_fp32` forces -every such param back to fp32 immediately before the load (the copy then lands -fp32 -> fp32). The checkpoint stores this tensor as fp32. - -The dense-vs-MoE split follows ``is_moe_layer`` (``first_k_dense_replace`` / -``moe_layer_freq``): early dense layers carry ``mlp.{gate,up,down}_proj`` into a -``ParallelGatedMLP``; MoE layers carry the expert-stacked + shared + gate params. -The routing here is layer-agnostic — a dense layer simply never emits -``mlp.experts.*`` / ``mlp.gate.*`` keys, and a MoE layer never emits a bare -``mlp.gate_proj``. - -Refs: HF key -> param authority is vLLM -``model_executor/models/deepseek_v2.py::DeepseekV2ForCausalLM.load_weights`` -(the ``stacked_params_mapping`` + per-expert ``expert_params_mapping`` there); -the fused ``w13``/``w2`` naming is vLLM's. - ----------------------------------------------------------------------------- -DESIGN NOTE — compressed-tensors INT4 / fp8 dequant-on-load (DEFERRED) ----------------------------------------------------------------------------- -The real Kimi-K2.7 checkpoint ships **compressed-tensors INT4** (fp8 variants -also exist); ``fused_experts`` is bf16/fp16-only and a full bf16 dequant of the -1T model is ~2 TB > 8xH200 (see the perf memo). So the real quantized checkpoint -is **out of scope for M5** (no checkpoint present, would not fit) — this loader is -the clean **bf16 path**, validated on a synthetic ``reduced()`` checkpoint. - -When the checkpoint + a quantized kernel land, dequant-on-load slots in **without -touching the routing above**, because ``load_hf_weights`` dispatches per parameter -through each param's ``weight_loader``: - - * compressed-tensors stores, per quantized tensor, a packed ``*.weight_packed`` - (INT4 nibbles / fp8 bytes) plus ``*.weight_scale`` (+ optional - ``*.weight_zero_point``) at a group granularity from the checkpoint's - ``quantization_config``. - * Hook point A (streaming): wrap the ``(name, tensor)`` iterator with a - dequantizer that consumes the ``weight_packed``/``weight_scale`` group, emits - a single bf16 ``*.weight`` (unpack nibble -> int -> ``(q - zp) * scale`` per - group), and drops the scale/zp/packed keys. Downstream routing is unchanged. - * Hook point B (per-param, memory-lean): keep the packed tensor in VRAM and - give the *destination* fused param a quant-aware ``weight_loader`` that stores - packed shards + scales (extend :class:`KimiSparseMoeBlock` to hold - ``gate_up_proj_packed`` / ``_scale``) and swap ``_dispatch`` for a quantized - grouped-GEMM. This is the only way to actually *serve* the 1T model and is - tracked as the top memory item in the perf backlog — a separate effort from - this DeepSeek-V3 port. - -Either hook is additive: the bf16 routing (remap + stacked rules) below is the -substrate both build on. +"""Kimi-K2.7 / DeepSeek-V3 weight loading. + +Maps an HF ``DeepseekV3ForCausalLM`` checkpoint onto the Kimi module tree via the +shared ``load_hf_weights`` machinery (name remap + stacked-shard rules), mirroring +``qwen3_omni_model.py``'s thinker remap/stacked params. + +- :func:`kimi_name_remapper`: strip a ``language_model.`` prefix if present (the + multimodal K2.7-Code repo carries it on its text weights); + ``shared_experts`` -> ``shared_expert``; tag per-routed-expert projections with + an ``__expert{i}__`` marker so one ``shard_id`` carries projection + expert slot. + Vision (``vision_tower.*`` / ``mm_projector.*``) and ``weight_shape`` keys fall + through and are dropped by the base loader's unknown-key skip. +- :func:`build_kimi_stacked_params`: fuse per-expert gate/up -> ``gate_up_proj`` + (w13) and down -> ``down_proj`` (w2); dense/shared gate+up -> ``gate_up_proj``. + Dense rules MUST come after the expert rules — ``_apply_stacked`` returns on + first match and a remapped expert key also contains ``.gate_proj``. +- Router bias (``e_score_correction_bias``) is forced fp32 before load so the + whole-model ``.to(bf16)`` cast can't downcast this fp32 selection bias. + +MLA loads strictly by name — the naive path keeps separate ``q_a_proj`` / +``kv_a_proj_with_mqa``, so no ``fused_qkv_a_proj`` fusion is needed. + +Compressed-tensors INT4 checkpoints: with a ``quant_config`` the stream is +dequantized to bf16 on load (see ``quantization.py``); routed experts can instead +stay packed (``packed_experts=True``) and dequantize inside the fused-expert +kernel. Both are additive — the remap + stacked rules are unchanged. + +Ref: HF key -> param authority is vLLM +``model_executor/models/deepseek_v2.py::DeepseekV2ForCausalLM.load_weights``. """ from __future__ import annotations import re from collections.abc import Iterable from pathlib import Path +from typing import TYPE_CHECKING import torch from torch import nn from mstar.model.loader.base import StackedParamRule -# HF checkpoint suffixes for the per-routed-expert projections. +if TYPE_CHECKING: + from mstar.model.kimi_k2_7.quantization import CompressedTensorsQuantConfig + +# HF suffixes for the per-routed-expert projections. The trailing alternation +# covers a native-bf16 ``.weight`` AND the compressed-tensors sub-keys +# (``.weight_packed`` / ``.weight_scale`` / ``.weight_zero_point``) so a +# packed-expert stream (which passes those sub-keys through raw) is remapped with +# its expert index preserved. A dequant-on-load stream only ever carries ``.weight``. _EXPERT_RE = re.compile( - r"(.*)\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$" + r"(.*)\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)" + r"\.(weight|weight_packed|weight_scale|weight_zero_point)$" ) +# Base-name matcher (no suffix) for the routed-expert weights kept packed. Used to +# build the ``keep_packed`` predicate handed to the dequant stream. +_EXPERT_BASE_RE = re.compile(r"\.experts\.\d+\.(gate_proj|up_proj|down_proj)$") + + +def _is_routed_expert_base(base: str) -> bool: + """True for a routed-expert weight base (``...experts..``). + + ``shared_experts`` does not match — there is no ``.experts..`` (the HF + key is ``mlp.shared_experts.gate_proj``, an underscore not a dotted index), so + the shared expert still dequantizes on load while the routed experts stay packed. + """ + return _EXPERT_BASE_RE.search(base) is not None + def kimi_name_remapper(name: str) -> str | None: """HF DeepSeek-V3 checkpoint key -> Kimi module param path. Returns ``None`` to drop a key (precomputed ``rotary_emb`` buffers). See the module docstring for the full mapping; MLA / norms / embed / lm_head are all - identity. + identity. Vision (``vision_tower.*`` / ``mm_projector.*``) and ``.weight_shape`` + sub-keys are left unmapped and fall through the base loader's unknown-key skip. """ if "rotary_emb" in name: return None + # Multimodal K2.7-Code text keys carry a ``language_model.`` prefix; strip it + # only-if-present (a bare ``model.*`` key is left unchanged). + if name.startswith("language_model."): + name = name[len("language_model."):] # HF names the shared expert plural; our module has one ``shared_expert``. name = name.replace(".shared_experts.", ".shared_expert.") - # Per-expert fusion marker so the stacked rules can pick up expert index. + # Per-expert fusion marker so the stacked rules can pick up expert index. The + # suffix (``weight`` for bf16, ``weight_packed``/``weight_scale`` for packed + # experts) is carried through so the packed vs bf16 stacked rules can route it. m = _EXPERT_RE.match(name) if m: - prefix, expert_idx, proj = m.groups() - return f"{prefix}.experts.{proj}.__expert{expert_idx}__.weight" + prefix, expert_idx, proj, suffix = m.groups() + return f"{prefix}.experts.{proj}.__expert{expert_idx}__.{suffix}" return name -def build_kimi_stacked_params(n_routed_experts: int) -> list[StackedParamRule]: +def build_kimi_stacked_params( + n_routed_experts: int, packed_experts: bool = False, +) -> list[StackedParamRule]: """Fused-shard routing for Kimi-K2.7 (mirrors the Qwen3-MoE thinker rules). - Per-expert ``gate``/``up`` -> ``experts.gate_up_proj`` (w13) and ``down`` -> - ``experts.down_proj`` (w2), then the dense/shared SwiGLU gate/up merge. - Expert rules precede the dense rules (first-match wins in ``_apply_stacked``). + ``packed_experts=False`` (native / dequantized bf16): per-expert ``gate``/``up`` + -> ``experts.gate_up_proj`` (w13) and ``down`` -> ``experts.down_proj`` (w2). + + ``packed_experts=True``: the per-expert ``.weight_packed`` / ``.weight_scale`` + sub-keys route to the FOUR packed params + (``experts.{gate_up_proj,down_proj}_{packed,scale}``), and the bf16 ``.weight`` + expert rules are OMITTED — their ``...__expert{i}__.weight`` source substring + would spuriously match ``...__expert{i}__.weight_packed`` (first-match wins). + + The dense/shared SwiGLU gate/up merge is appended last in both cases (expert + rules precede it so ``.gate_proj`` inside an expert key can't hijack it). """ rules: list[StackedParamRule] = [] for i in range(n_routed_experts): - rules.append(StackedParamRule( - target_suffix=".experts.gate_up_proj", - source_suffix=f".experts.gate_proj.__expert{i}__.weight", - shard_id=f"gate:{i}", - )) - rules.append(StackedParamRule( - target_suffix=".experts.gate_up_proj", - source_suffix=f".experts.up_proj.__expert{i}__.weight", - shard_id=f"up:{i}", - )) - rules.append(StackedParamRule( - target_suffix=".experts.down_proj", - source_suffix=f".experts.down_proj.__expert{i}__.weight", - shard_id=f"down:{i}", - )) + if packed_experts: + for proj, sid in (("gate_proj", f"gate:{i}"), ("up_proj", f"up:{i}")): + rules.append(StackedParamRule( + target_suffix=".experts.gate_up_proj_packed", + source_suffix=f".experts.{proj}.__expert{i}__.weight_packed", + shard_id=sid, + )) + rules.append(StackedParamRule( + target_suffix=".experts.gate_up_proj_scale", + source_suffix=f".experts.{proj}.__expert{i}__.weight_scale", + shard_id=sid, + )) + rules.append(StackedParamRule( + target_suffix=".experts.down_proj_packed", + source_suffix=f".experts.down_proj.__expert{i}__.weight_packed", + shard_id=f"down:{i}", + )) + rules.append(StackedParamRule( + target_suffix=".experts.down_proj_scale", + source_suffix=f".experts.down_proj.__expert{i}__.weight_scale", + shard_id=f"down:{i}", + )) + else: + rules.append(StackedParamRule( + target_suffix=".experts.gate_up_proj", + source_suffix=f".experts.gate_proj.__expert{i}__.weight", + shard_id=f"gate:{i}", + )) + rules.append(StackedParamRule( + target_suffix=".experts.gate_up_proj", + source_suffix=f".experts.up_proj.__expert{i}__.weight", + shard_id=f"up:{i}", + )) + rules.append(StackedParamRule( + target_suffix=".experts.down_proj", + source_suffix=f".experts.down_proj.__expert{i}__.weight", + shard_id=f"down:{i}", + )) # Dense MLP + shared-expert gate/up fusion — AFTER the expert rules. rules.append(StackedParamRule(".gate_up_proj", ".gate_proj", 0)) rules.append(StackedParamRule(".gate_up_proj", ".up_proj", 1)) @@ -169,21 +174,46 @@ def load_kimi_hf_weights( module: nn.Module, weights: Iterable[tuple[str, torch.Tensor]], n_routed_experts: int, + quant_config: "CompressedTensorsQuantConfig | None" = None, + packed_experts: bool = False, ) -> set[str]: """Load an HF DeepSeek-V3 weight stream into ``module``. - Thin wrapper: restore the fp32 router bias, then dispatch through - ``load_hf_weights`` with the Kimi remap + stacked rules. Returns the set of - param paths that received a tensor (callers can diff against - ``named_parameters()`` to assert completeness). + Thin wrapper: restore the fp32 router bias, optionally wrap the stream with the + dequant-on-load parser (when ``quant_config`` is set — the checkpoint is + compressed-tensors quantized), then dispatch through ``load_hf_weights`` with + the Kimi remap + stacked rules. The dequant wrapper emits bf16 ``*.weight`` + keys, so the remap + stacked rules see the same stream as a native-bf16 + checkpoint. Returns the set of param paths that received a tensor (callers can + diff against ``named_parameters()`` to assert completeness). + + ``packed_experts=True``: the routed experts stay PACKED. A ``keep_packed`` + predicate matching routed-expert bases is handed to the dequant stream so those + sub-keys pass through raw (int32 + scale), and the stacked rules route them to + the packed params; every other quantized weight (MLA, dense FFN, shared expert) + still dequantizes to bf16. Requires ``quant_config``. """ from mstar.model.loader import load_hf_weights + if quant_config is not None: + from mstar.model.kimi_k2_7.quantization import ( + dequant_compressed_tensors_stream, + ) + + keep_packed = _is_routed_expert_base if packed_experts else None + weights = dequant_compressed_tensors_stream( + weights, quant_config, keep_packed=keep_packed, + ) + elif packed_experts: + raise ValueError("packed_experts=True requires a quant_config") + restore_router_bias_fp32(module) return load_hf_weights( module, weights, - stacked_params=build_kimi_stacked_params(n_routed_experts), + stacked_params=build_kimi_stacked_params( + n_routed_experts, packed_experts=packed_experts, + ), name_remapper=kimi_name_remapper, ) diff --git a/mstar/utils/fused_moe/kernels.py b/mstar/utils/fused_moe/kernels.py index e31c6daae..a8fb61993 100644 --- a/mstar/utils/fused_moe/kernels.py +++ b/mstar/utils/fused_moe/kernels.py @@ -138,6 +138,147 @@ def fused_moe_kernel( tl.store(c_ptrs, accumulator, mask=c_mask) +@triton.jit +def fused_moe_kernel_w4a16( + # Pointers + a_ptr, + b_ptr, + c_ptr, + b_scale_ptr, + b_zp_ptr, + topk_weights_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_tokens_post_padded_ptr, + # Matrix dimensions (K is the LOGICAL contraction dim, not the packed width) + N, + K, + EM, + num_valid_tokens, + # Strides + stride_am, + stride_ak, + stride_be, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_bse, + stride_bsk, + stride_bsn, + stride_bze, + stride_bzk, + stride_bzn, + # Block sizes (compile-time) + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + MUL_ROUTED_WEIGHT: tl.constexpr, + top_k: tl.constexpr, + compute_type: tl.constexpr, + group_size: tl.constexpr, + PACK_FACTOR: tl.constexpr, + HAS_ZP: tl.constexpr, + even_Ks: tl.constexpr, +): + """Compute one ``[BLOCK_SIZE_M, BLOCK_SIZE_N]`` output tile from PACKED weights. + + Identical control flow to :func:`fused_moe_kernel`; the only change is the + B load in the K-loop. ``b_ptr`` addresses an int32 tensor of shape + ``(E, N, K // PACK_FACTOR)`` where ``PACK_FACTOR`` INT4 nibbles are packed + low-order-first along the (logical) K axis into each int32. For each K tile we + read the containing int32s, shift out the right nibble, offset-binary subtract + (symmetric: ``- 8``; asymmetric: ``- b_zp``), and scale by the per-``group_size`` + ``b_scale``, all in fp32, then cast to ``compute_type`` for the ``tl.dot``. + + Requires ``BLOCK_SIZE_K % PACK_FACTOR == 0`` and ``BLOCK_SIZE_K % group_size + == 0`` (enforced by :func:`get_default_config`) so a K tile spans a whole + number of int32s and never straddles a group boundary mid-int32. + """ + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr) + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: + return + + offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64) + token_mask = offs_token < num_valid_tokens + + off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak) + # Packed B: the int32 holding logical-K index ``kk`` is at ``kk // PACK_FACTOR`` + # along the last (packed) axis; the nibble sits at ``(kk % PACK_FACTOR) * 4``. + b_ptrs = ( + b_ptr + off_experts * stride_be + + (offs_k[:, None] // PACK_FACTOR) * stride_bk + + offs_bn[None, :] * stride_bn + ) + b_shifter = (offs_k[:, None] % PACK_FACTOR) * 4 + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for k_start in range(0, K, BLOCK_SIZE_K): + # Per-group scale index along the logical K axis (same for every int32 in + # a group). Recomputed each tile from the global-K position. + offs_ks = (offs_k[:, None] + k_start) // group_size + b_scale_ptrs = b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn + offs_ks * stride_bsk + if even_Ks: + a = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0) + b_packed = tl.load(b_ptrs) + b_scale = tl.load(b_scale_ptrs).to(tl.float32) + else: + k_mask = offs_k[:, None] < K - k_start + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < K - k_start), + other=0.0, + ) + b_packed = tl.load(b_ptrs, mask=k_mask, other=0) + b_scale = tl.load(b_scale_ptrs, mask=k_mask, other=1.0).to(tl.float32) + # Extract the nibble. ``>>`` on int32 is arithmetic, but ``& 0xF`` masks + # the sign-extended high bits, so the top nibble (container bit 31 set) is + # exact for all PACK_FACTOR positions. + b_nib = ((b_packed >> b_shifter) & 0xF).to(tl.float32) + if HAS_ZP: + # Asymmetric extension (never exercised by Kimi — symmetric only). The + # zero point is stored one-per-group, unpacked, mirroring ``b_scale``. + b_zp_ptrs = b_zp_ptr + off_experts * stride_bze + offs_bn[None, :] * stride_bzn + offs_ks * stride_bzk + if even_Ks: + b_zp = tl.load(b_zp_ptrs).to(tl.float32) + else: + b_zp = tl.load(b_zp_ptrs, mask=offs_k[:, None] < K - k_start, other=0.0).to(tl.float32) + b = ((b_nib - b_zp) * b_scale).to(compute_type) + else: + b = ((b_nib - 8.0) * b_scale).to(compute_type) + accumulator += tl.dot(a, b) + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += (BLOCK_SIZE_K // PACK_FACTOR) * stride_bk + + if MUL_ROUTED_WEIGHT: + moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0) + accumulator = accumulator * moe_weight[:, None] + + accumulator = accumulator.to(compute_type) + + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + + def invoke_fused_moe_kernel( A: torch.Tensor, B: torch.Tensor, @@ -218,6 +359,93 @@ def grid(META): ) +def invoke_fused_moe_kernel_w4a16( + A: torch.Tensor, + B_packed: torch.Tensor, + C: torch.Tensor, + B_scale: torch.Tensor, + B_zp: torch.Tensor | None, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + mul_routed_weight: bool, + top_k: int, + config: Dict[str, Any], + compute_type: tl.dtype, + K: int, + pack_factor: int, + group_size: int, +) -> None: + """Launch :func:`fused_moe_kernel_w4a16` (packed-INT4 grouped GEMM). + + Mirrors :func:`invoke_fused_moe_kernel` with two differences that the packed + layout forces: + + * ``K`` is the LOGICAL contraction dim and is passed explicitly — it CANNOT be + derived from ``B_packed.shape[2]`` (that is ``K // pack_factor``). + * ``B_scale`` (shape ``(E, N, K // group_size)``) rides alongside; its strides + are passed ``stride(0), stride(2), stride(1)`` to match the kernel's + ``(bse, bsk, bsn)`` order, exactly like the weight strides. + + ``B_zp`` is optional (symmetric INT4 has no zero point). When ``None`` the + kernel's ``HAS_ZP`` is off and a stand-in tensor (``B_scale``) is passed so the + dead zp-stride args stay valid; nothing is dereferenced. + """ + assert topk_weights.stride(1) == 1 + assert sorted_token_ids.stride(0) == 1 + + N = B_packed.shape[1] + + def grid(META): + return ( + triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"]) + * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ) + + even_Ks = (K % config["BLOCK_SIZE_K"]) == 0 + has_zp = B_zp is not None + zp = B_zp if has_zp else B_scale # stand-in; every zp load is gated by HAS_ZP + + fused_moe_kernel_w4a16[grid]( + A, + B_packed, + C, + B_scale, + zp, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + N, + K, + sorted_token_ids.shape[0], + topk_ids.numel(), + A.stride(0), + A.stride(1), + B_packed.stride(0), + B_packed.stride(2), + B_packed.stride(1), + C.stride(-2), + C.stride(-1), + B_scale.stride(0), + B_scale.stride(2), + B_scale.stride(1), + zp.stride(0), + zp.stride(2), + zp.stride(1), + MUL_ROUTED_WEIGHT=mul_routed_weight, + top_k=top_k, + compute_type=compute_type, + group_size=group_size, + PACK_FACTOR=pack_factor, + HAS_ZP=has_zp, + even_Ks=even_Ks, + **config, + ) + + # --------------------------------------------------------------------------- # Activation (SwiGLU / GeGLU) kernel # --------------------------------------------------------------------------- @@ -398,23 +626,47 @@ def moe_sum_reduce_triton( # --------------------------------------------------------------------------- -def get_default_config(M: int, E: int, N: int, K: int, top_k: int) -> Dict[str, int]: +def get_default_config( + M: int, E: int, N: int, K: int, top_k: int, group_size: int | None = None, +) -> Dict[str, int]: """Pick Triton tile sizes based on problem shape. Mirrors sglang's ``get_default_config`` for the unquantized path. For decode batch sizes (``M`` on the order of 1--64) we always fall into the ``M <= E`` branch since Qwen3-Omni has ``E == 128``. + + ``group_size`` (set only on the W4A16 path) does NOT change the bf16 result: + when ``None`` the returned config is byte-for-byte the historical one. When + set, ``BLOCK_SIZE_K`` is clamped down (halving) until it is a multiple of both + ``pack_factor`` (8 for INT4) and ``group_size`` — the divisibility the packed + kernel needs so a K tile spans whole int32s and whole groups. Kimi + (``group_size=32``; configs 64/32) already complies, so the clamp is a no-op. """ if M <= E: - return { + config = { "BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 1, } - return { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 64, - "BLOCK_SIZE_K": 32, - "GROUP_SIZE_M": 8, - } + else: + config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + } + if group_size is not None: + pack_factor = 8 # INT4: 32 // 4 + bk = config["BLOCK_SIZE_K"] + while bk % pack_factor != 0 or bk % group_size != 0: + bk //= 2 + if bk < pack_factor: + raise ValueError( + f"cannot pick a BLOCK_SIZE_K divisible by pack_factor={pack_factor} " + f"and group_size={group_size}; got down to {bk}" + ) + config["BLOCK_SIZE_K"] = bk + assert config["BLOCK_SIZE_K"] % pack_factor == 0 + assert config["BLOCK_SIZE_K"] % group_size == 0 + return config diff --git a/mstar/utils/fused_moe/runner.py b/mstar/utils/fused_moe/runner.py index 0926fa49d..66d98bd16 100644 --- a/mstar/utils/fused_moe/runner.py +++ b/mstar/utils/fused_moe/runner.py @@ -16,6 +16,7 @@ act_and_mul_triton, get_default_config, invoke_fused_moe_kernel, + invoke_fused_moe_kernel_w4a16, moe_sum_reduce_triton, ) @@ -36,6 +37,12 @@ def fused_experts( topk_ids: torch.Tensor, activation: str = "silu", reduce_results: bool = True, + w1_scale: torch.Tensor | None = None, + w2_scale: torch.Tensor | None = None, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, + group_size: int | None = None, + pack_factor: int | None = None, ) -> torch.Tensor: """Grouped-GEMM Triton MoE dispatch. @@ -48,10 +55,13 @@ def fused_experts( ``(num_experts, 2 * moe_intermediate_size, hidden)``. Matches the ``experts.gate_up_proj`` parameter the WeightConverter already produces in :mod:`mstar.model.qwen3_omni.qwen3_omni_model`. + On the W4A16 path (``w1_scale`` set) this is instead the PACKED int32 + tensor ``(num_experts, 2 * moe_intermediate_size, hidden // pack_factor)``. w2 : torch.Tensor Down projection weights, shape ``(num_experts, hidden, moe_intermediate_size)``. Matches - ``experts.down_proj``. + ``experts.down_proj``. On the W4A16 path this is the packed int32 tensor + ``(num_experts, hidden, moe_intermediate_size // pack_factor)``. topk_weights : torch.Tensor ``(tokens, top_k)``, routing probabilities (possibly renormalized). Dtype matches ``hidden_states``. @@ -64,6 +74,15 @@ def fused_experts( ``(tokens, hidden)``. If False, skip the sum-reduce and return ``(tokens, top_k, hidden)`` — the caller is responsible for the reduce (e.g. after an all-reduce for TP). + w1_scale, w2_scale : torch.Tensor | None + W4A16 group scales, shapes ``(E, 2*inter, hidden//group_size)`` + and ``(E, hidden, inter//group_size)``. When BOTH are ``None`` (default) + the bf16 path runs byte-for-byte as before; when set, ``w1``/``w2`` are + packed int32 and the in-kernel dequant path runs. + w1_zp, w2_zp : torch.Tensor | None + Optional asymmetric zero points (unused for Kimi's symmetric INT4). + group_size, pack_factor : int | None + Group granularity and INT4 pack factor (8); required on the W4A16 path. Returns ------- @@ -72,25 +91,50 @@ def fused_experts( ``(tokens, top_k, hidden)``. """ assert hidden_states.is_contiguous(), "hidden_states must be contiguous" - assert w1.is_contiguous(), "w1 must be contiguous" - assert w2.is_contiguous(), "w2 must be contiguous" assert hidden_states.dim() == 2 assert topk_weights.shape == topk_ids.shape + # Activations stay bf16/fp16 on both paths — only the WEIGHTS are quantized. assert hidden_states.dtype in (torch.bfloat16, torch.float16) + quantized = w1_scale is not None num_tokens, hidden = hidden_states.shape - E, two_inter, k_in = w1.shape - assert k_in == hidden, f"w1 last dim {k_in} != hidden {hidden}" - _, w2_hidden, inter = w2.shape - assert w2_hidden == hidden, f"w2 dim[1] {w2_hidden} != hidden {hidden}" - assert two_inter == 2 * inter, f"w1 dim[1] {two_inter} != 2 * w2 dim[2] {2 * inter}" + if quantized: + assert pack_factor is not None and group_size is not None, ( + "W4A16 path requires pack_factor and group_size" + ) + assert w2_scale is not None, "W4A16 path requires both w1_scale and w2_scale" + assert w1.dtype == torch.int32 and w2.dtype == torch.int32, ( + "W4A16 path expects packed int32 weights" + ) + assert w1.is_contiguous(), "w1 (packed) must be contiguous" + assert w2.is_contiguous(), "w2 (packed) must be contiguous" + E, two_inter, k1_packed = w1.shape + assert k1_packed == hidden // pack_factor, ( + f"w1 packed last dim {k1_packed} != hidden//pack_factor {hidden // pack_factor}" + ) + _, w2_hidden, k2_packed = w2.shape + assert w2_hidden == hidden, f"w2 dim[1] {w2_hidden} != hidden {hidden}" + inter = two_inter // 2 + assert k2_packed == inter // pack_factor, ( + f"w2 packed last dim {k2_packed} != inter//pack_factor {inter // pack_factor}" + ) + else: + assert w1.is_contiguous(), "w1 must be contiguous" + assert w2.is_contiguous(), "w2 must be contiguous" + E, two_inter, k_in = w1.shape + assert k_in == hidden, f"w1 last dim {k_in} != hidden {hidden}" + _, w2_hidden, inter = w2.shape + assert w2_hidden == hidden, f"w2 dim[1] {w2_hidden} != hidden {hidden}" + assert two_inter == 2 * inter, f"w1 dim[1] {two_inter} != 2 * w2 dim[2] {2 * inter}" top_k = topk_ids.shape[1] # moe_align_block_size expects int32; torch.topk returns int64. topk_ids = topk_ids.to(torch.int32).contiguous() topk_weights = topk_weights.contiguous() - config = get_default_config(M=num_tokens, E=E, N=two_inter, K=hidden, top_k=top_k) + config = get_default_config( + M=num_tokens, E=E, N=two_inter, K=hidden, top_k=top_k, group_size=group_size, + ) compute_type = _tl_compute_type(hidden_states.dtype) # 1. Token permute + per-expert block alignment. @@ -115,20 +159,42 @@ def fused_experts( ) # 3. Gate+up GEMM: cache1[slot] = hidden[slot // top_k] @ w1[expert].T - invoke_fused_moe_kernel( - A=hidden_states, - B=w1, - C=cache1, - topk_weights=topk_weights, - topk_ids=topk_ids, - sorted_token_ids=sorted_token_ids, - expert_ids=expert_ids, - num_tokens_post_padded=num_tokens_post_padded, - mul_routed_weight=False, - top_k=top_k, - config=config, - compute_type=compute_type, - ) + # (W4A16: w1 packed int32, dequantized in-kernel; GEMM-1 contracts over hidden.) + if quantized: + invoke_fused_moe_kernel_w4a16( + A=hidden_states, + B_packed=w1, + C=cache1, + B_scale=w1_scale, + B_zp=w1_zp, + topk_weights=topk_weights, + topk_ids=topk_ids, + sorted_token_ids=sorted_token_ids, + expert_ids=expert_ids, + num_tokens_post_padded=num_tokens_post_padded, + mul_routed_weight=False, + top_k=top_k, + config=config, + compute_type=compute_type, + K=hidden, + pack_factor=pack_factor, + group_size=group_size, + ) + else: + invoke_fused_moe_kernel( + A=hidden_states, + B=w1, + C=cache1, + topk_weights=topk_weights, + topk_ids=topk_ids, + sorted_token_ids=sorted_token_ids, + expert_ids=expert_ids, + num_tokens_post_padded=num_tokens_post_padded, + mul_routed_weight=False, + top_k=top_k, + config=config, + compute_type=compute_type, + ) # 4. SwiGLU: cache2[slot] = silu(gate) * up act_and_mul_triton(cache1, cache2, activation=activation) @@ -136,20 +202,42 @@ def fused_experts( # 5. Down GEMM (weighted): cache3[slot] = topk_weight[slot] * (cache2[slot] @ w2[expert].T) # top_k=1 for this GEMM so the kernel's offs_token // top_k is identity # -- it reads cache2 rows directly instead of the (slot // top_k)-th source row. - invoke_fused_moe_kernel( - A=cache2, - B=w2, - C=cache3.view(m_topk, hidden), - topk_weights=topk_weights, - topk_ids=topk_ids, - sorted_token_ids=sorted_token_ids, - expert_ids=expert_ids, - num_tokens_post_padded=num_tokens_post_padded, - mul_routed_weight=True, - top_k=1, - config=config, - compute_type=compute_type, - ) + # (W4A16: w2 packed int32; GEMM-2 contracts over the intermediate dim.) + if quantized: + invoke_fused_moe_kernel_w4a16( + A=cache2, + B_packed=w2, + C=cache3.view(m_topk, hidden), + B_scale=w2_scale, + B_zp=w2_zp, + topk_weights=topk_weights, + topk_ids=topk_ids, + sorted_token_ids=sorted_token_ids, + expert_ids=expert_ids, + num_tokens_post_padded=num_tokens_post_padded, + mul_routed_weight=True, + top_k=1, + config=config, + compute_type=compute_type, + K=inter, + pack_factor=pack_factor, + group_size=group_size, + ) + else: + invoke_fused_moe_kernel( + A=cache2, + B=w2, + C=cache3.view(m_topk, hidden), + topk_weights=topk_weights, + topk_ids=topk_ids, + sorted_token_ids=sorted_token_ids, + expert_ids=expert_ids, + num_tokens_post_padded=num_tokens_post_padded, + mul_routed_weight=True, + top_k=1, + config=config, + compute_type=compute_type, + ) # 6. Sum over the top-k slots -> (tokens, hidden). if not reduce_results: diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 97be3ed33..204edc912 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -985,7 +985,14 @@ def _send_outputs( request_id, outputs.persist ) - if outputs.new_token_outputs: + # Only the tp-leader computes new-token counts: the conductor consumes + # ``new_token_counts`` exclusively from ``is_first_tp_rank`` messages + # (see conductor.py — ``if body.is_first_tp_rank``), and for a replicated + # (non-persisted) EMIT_TO_CLIENT + conductor_new_token edge the token + # tensor is only kept alive on rank 0. Without this gate the 7 non-leader + # ranks call get_tensor() on an already-dereferenced/GC'd uuid and crash + # with KeyError. + if outputs.new_token_outputs and outputs.is_first_tp_rank: name_to_count: dict[str, int] = {} for signal in outputs.new_token_outputs: if signal.name in name_to_count: diff --git a/test/integration/test_kimi_moe_inkernel_dequant.py b/test/integration/test_kimi_moe_inkernel_dequant.py new file mode 100644 index 000000000..eaa966410 --- /dev/null +++ b/test/integration/test_kimi_moe_inkernel_dequant.py @@ -0,0 +1,123 @@ +"""GPU kernel golden: W4A16 in-kernel INT4 dequant vs the bf16 fused-expert GEMM. + +The in-kernel dequant path ships a SEPARATE ``fused_moe_kernel_w4a16`` that keeps the routed +experts packed in VRAM and dequantizes each K tile in registers before the dot. +Its correctness invariant is exact: the nibble ``(q - 8) * scale`` cast to bf16 is +*the same value* the bf16 path feeds to ``tl.dot`` after a pre-dequant, and with +the same tile config the two accumulate in the same order — so the packed path +must match the bf16 path on the SAME dequantized weights to a tight tolerance. + +This is the cheapest level that catches a kernel bug (packed-K stride, nibble +shifter, group-scale index, the top-nibble sign case) without a full model: + + 1. random bf16 experts ``w1 (E, 2I, H)`` / ``w2 (E, H, I)``; + 2. ``fake_quantize_weight`` each expert to ``(packed, bf16 scale, deq_bf16)``, + with ``scale_dtype=bfloat16`` so the packed-param scale and the bf16-path + weight dequantize from the identical scale; + 3. assert ``fused_experts(x, w1_packed, w2_packed, w1_scale=, w2_scale=, ...)`` + == ``fused_experts(x, w1_deq, w2_deq)`` (the bf16 path). + +Includes an expert whose packing sets container bit 31 (top nibble >= 8), proving +the arithmetic-shift + ``& 0xF`` mask recovers it. + +Run: pytest test/integration/test_kimi_moe_inkernel_dequant.py -v +""" +import pytest +import torch + +from mstar.model.kimi_k2_7._testing import fake_quantize_weight +from mstar.model.kimi_k2_7.quantization import unpack_int32 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="W4A16 fused-expert kernel golden needs a GPU", +) + +DEVICE = "cuda" +GROUP_SIZE = 32 +PACK_FACTOR = 8 + + +def _quantize_stack(weight): + """Fake-quantize a stacked ``(E, N, K)`` weight, returning packed/scale/deq. + + Each expert is quantized independently (matching a per-Linear checkpoint); + the bf16 scale is what a real compressed-tensors checkpoint stores, so the + returned ``deq`` is bit-for-bit what the packed param dequantizes to. + """ + E, N, K = weight.shape + packed = torch.empty((E, N, K // PACK_FACTOR), dtype=torch.int32, device=DEVICE) + scale = torch.empty((E, N, K // GROUP_SIZE), dtype=torch.bfloat16, device=DEVICE) + deq = torch.empty((E, N, K), dtype=torch.bfloat16, device=DEVICE) + for e in range(E): + p, s, d = fake_quantize_weight( + weight[e], num_bits=4, group_size=GROUP_SIZE, symmetric=True, + scale_dtype=torch.bfloat16, + ) + packed[e], scale[e], deq[e] = p.to(DEVICE), s.to(DEVICE), d.to(DEVICE) + return packed, scale, deq + + +def _random_topk(num_tokens, E, top_k): + logits = torch.randn(num_tokens, E, device=DEVICE) + weights, ids = torch.topk(logits.softmax(-1), top_k, dim=-1) + weights = weights / weights.sum(-1, keepdim=True) + return weights.to(torch.bfloat16), ids + + +@pytest.mark.parametrize("num_tokens", [8, 3]) # M > E and M <= E branches +def test_w4a16_matches_bf16_on_same_dequant(num_tokens): + from mstar.utils.fused_moe.runner import fused_experts + + torch.manual_seed(0) + E, H, I, top_k = 4, 128, 64, 2 + # Slightly wide init so per-group amax spans the full nibble range and some + # top nibbles land >= 8 (container bit 31 set) — the sign-mask path. + w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16) + w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16) + + w1_packed, w1_scale, w1_deq = _quantize_stack(w1) + w2_packed, w2_scale, w2_deq = _quantize_stack(w2) + + # Guard: the packing really exercises the negative-container / top-nibble>=8 + # case (else the sign-extension mask would be untested). + assert (w1_packed < 0).any(), "no int32 with bit 31 set — top-nibble path untested" + top_nibbles = unpack_int32(w1_packed.cpu(), num_bits=4)[..., PACK_FACTOR - 1 :: PACK_FACTOR] + assert (top_nibbles >= 8).any() + + x = (torch.randn(num_tokens, H, device=DEVICE) * 0.5).to(torch.bfloat16) + topk_weights, topk_ids = _random_topk(num_tokens, E, top_k) + + out_quant = fused_experts( + x, w1_packed, w2_packed, topk_weights, topk_ids, + w1_scale=w1_scale, w2_scale=w2_scale, group_size=GROUP_SIZE, pack_factor=PACK_FACTOR, + ) + out_bf16 = fused_experts(x, w1_deq, w2_deq, topk_weights, topk_ids) + + assert out_quant.shape == (num_tokens, H) + assert out_quant.dtype == torch.bfloat16 + torch.testing.assert_close(out_quant, out_bf16, rtol=1e-2, atol=1e-2) + + +def test_w4a16_reduce_results_false_shape(): + """``reduce_results=False`` returns the per-slot (tokens, top_k, hidden) tensor + the TP path all-reduces before the top-k sum — exercise it on the packed path.""" + from mstar.utils.fused_moe.runner import fused_experts + + torch.manual_seed(1) + E, H, I, top_k, num_tokens = 4, 128, 64, 2, 6 + w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16) + w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16) + w1_packed, w1_scale, w1_deq = _quantize_stack(w1) + w2_packed, w2_scale, w2_deq = _quantize_stack(w2) + x = (torch.randn(num_tokens, H, device=DEVICE) * 0.5).to(torch.bfloat16) + topk_weights, topk_ids = _random_topk(num_tokens, E, top_k) + + got = fused_experts( + x, w1_packed, w2_packed, topk_weights, topk_ids, + w1_scale=w1_scale, w2_scale=w2_scale, group_size=GROUP_SIZE, pack_factor=PACK_FACTOR, + reduce_results=False, + ) + exp = fused_experts(x, w1_deq, w2_deq, topk_weights, topk_ids, reduce_results=False) + assert got.shape == (num_tokens, top_k, H) + torch.testing.assert_close(got, exp, rtol=1e-2, atol=1e-2) diff --git a/test/integration/test_kimi_quant_inkernel_weight_loading.py b/test/integration/test_kimi_quant_inkernel_weight_loading.py new file mode 100644 index 000000000..5c730c8ac --- /dev/null +++ b/test/integration/test_kimi_quant_inkernel_weight_loading.py @@ -0,0 +1,376 @@ +"""Golden: packed experts + in-kernel INT4 dequant for Kimi-K2.7. + +Where the dequant-on-load golden (``test_kimi_quant_weight_loading.py``) +dequantizes every quantized weight to bf16 on load, in-kernel dequant keeps the +ROUTED EXPERTS packed int32 in VRAM and dequantizes each tile inside +``fused_moe_kernel_w4a16``; MLA / dense-FFN / shared-expert weights still dequant +on load. This test loads BOTH models from the SAME synthetic compressed-tensors +checkpoint and pins the packed-expert behavior against the dequant-on-load model: + + 1. build a reference (bf16 experts) and fake-quantize every eligible weight in + place, serializing an HF compressed-tensors checkpoint (``weight_packed`` + + ``weight_scale`` for quantized weights, plain ``weight`` for the ignore set); + 2. load ``model_a`` with ``reduced_quantized()`` (dequant-on-load — experts dequant + to bf16) and ``model_b`` with ``reduced_quantized_inkernel()`` (packed experts — + experts stay packed) from that one checkpoint; + 3. assert (a) completeness — ``model_b``'s loaded set equals its + ``named_parameters()``, now carrying ``*_packed`` / ``*_scale`` and NOT the + bf16 fused expert params; (b) the packed params survived the whole-model + ``.to(bf16)`` cast as int32 (the downcast-exemption guard); (c) router bias + fp32, no stray buffers; (d) ``model_b``'s full forward matches ``model_a``'s + within a loose bf16 tolerance (dequant-on-load and in-kernel dequant differ + only in accumulation order — both dequant to the identical bf16 weights). + +A tp=2 packed-MoE-block simulation (``test_packed_moe_block_tp2_sim_matches_tp1``) +separately proves the packed per-rank weight_loaders' column/row slicing and the +``_dispatch_packed_experts`` all-reduce path (shard_inter=32 satisfies the pack / +group divisibility asserts). + +Run: pytest test/integration/test_kimi_quant_inkernel_weight_loading.py -v +""" +import pytest +import torch + +from mstar.distributed.communication import CommGroup +from mstar.model.kimi_k2_7._testing import fake_quantize_weight +from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM +from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock +from mstar.model.kimi_k2_7.config import KimiK2Config +from mstar.model.loader import load_weights as driver_load_weights + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Kimi packed-expert weight-loading golden needs a GPU (W4A16 fused expert GEMM)", +) + +DEVICE = "cuda" +DTYPE = torch.bfloat16 + + +# -------------------------------------------------------------------------- +# Mock paged cache (causal SDPA at 1/sqrt(head_dim)) — same as the dequant-on-load golden. +# -------------------------------------------------------------------------- + +def _sdpa_causal(q, k, v, scale): + qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) + T = q.shape[0] + causal = torch.triu(torch.full((T, T), float("-inf"), device=q.device), diagonal=1) + attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1) + return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) + + +class _MockMLACache: + def __init__(self, head_dim): + self.scale = head_dim ** -0.5 + + def set_layer_idx(self, _i): + pass + + def advance_seq_lens(self, *_a, **_k): + pass + + def run_attention(self, q, k, v): + return _sdpa_causal(q, k, v, self.scale) + + +# -------------------------------------------------------------------------- +# Random weight init + fake-quant serialization (mirrors the dequant-on-load golden). +# -------------------------------------------------------------------------- + +def _fill_layer(layer, cfg): + a = layer.self_attn + for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): + lin.weight.data.normal_(0, 0.03) + for norm in (a.q_a_layernorm, a.kv_a_layernorm): + norm.weight.data.normal_(1.0, 0.02) + layer.input_layernorm.weight.data.normal_(1.0, 0.02) + layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02) + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + mlp.gate.weight.data.normal_(0, 1) + mlp.gate.e_score_correction_bias.data = torch.randn( + cfg.n_routed_experts, device=DEVICE, dtype=torch.float32) + # Wide expert init so per-group amax spans the nibble range (exercises the + # top-nibble sign path in the kernel) — the dequant-on-load build has bf16 experts. + mlp.experts.gate_up_proj.data.normal_(0, 0.2) + mlp.experts.down_proj.data.normal_(0, 0.2) + mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05) + else: + mlp.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.down_proj.weight.data.normal_(0, 0.05) + + +def _build_reference(cfg): + """dequant-on-load config (bf16 fused experts) — the source the checkpoint is cut from.""" + model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=DTYPE) + model.model.embed_tokens.weight.data.normal_(0, 0.05) + model.model.norm.weight.data.normal_(1.0, 0.02) + model.lm_head.weight.data.normal_(0, 0.02) + for layer in model.model.layers: + _fill_layer(layer, cfg) + model.requires_grad_(False) + return model.eval() + + +def _keep_bf16(key): + """Weights a real compressed-tensors Kimi checkpoint leaves in bf16.""" + if key == "lm_head.weight": + return True + return ( + key.endswith("norm.weight") + or key.endswith("mlp.gate.weight") + or key.endswith("e_score_correction_bias") + or key.endswith("embed_tokens.weight") + ) + + +def _emit(sd, key, view, quant_cfg): + """Quantize (and write the dequant back into ``view``) when eligible, else store + the plain bf16 tensor — 2-D + input dim divisible by group_size + not bf16-kept.""" + eligible = ( + not _keep_bf16(key) + and view.dim() == 2 + and view.shape[-1] % quant_cfg.group_size == 0 + ) + if not eligible: + sd[key] = view + return + packed, scale, deq = fake_quantize_weight( + view, num_bits=quant_cfg.num_bits, group_size=quant_cfg.group_size, + symmetric=quant_cfg.symmetric, scale_dtype=DTYPE, + ) + base = key[: -len(".weight")] + sd[base + ".weight_packed"] = packed + sd[base + ".weight_scale"] = scale + view.copy_(deq.to(view.dtype)) + + +def _hf_quant_checkpoint(model, cfg, quant_cfg): + inter = cfg.intermediate_size + moe_inter = cfg.moe_intermediate_size + shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts + m = model.model + sd = {} + _emit(sd, "model.embed_tokens.weight", m.embed_tokens.weight, quant_cfg) + for i, layer in enumerate(m.layers): + p = f"model.layers.{i}." + a = layer.self_attn + _emit(sd, p + "self_attn.q_a_proj.weight", a.q_a_proj.weight, quant_cfg) + _emit(sd, p + "self_attn.q_a_layernorm.weight", a.q_a_layernorm.weight, quant_cfg) + _emit(sd, p + "self_attn.q_b_proj.weight", a.q_b_proj.weight, quant_cfg) + _emit(sd, p + "self_attn.kv_a_proj_with_mqa.weight", + a.kv_a_proj_with_mqa.weight, quant_cfg) + _emit(sd, p + "self_attn.kv_a_layernorm.weight", a.kv_a_layernorm.weight, quant_cfg) + _emit(sd, p + "self_attn.kv_b_proj.weight", a.kv_b_proj.weight, quant_cfg) + _emit(sd, p + "self_attn.o_proj.weight", a.o_proj.weight, quant_cfg) + _emit(sd, p + "input_layernorm.weight", layer.input_layernorm.weight, quant_cfg) + _emit(sd, p + "post_attention_layernorm.weight", + layer.post_attention_layernorm.weight, quant_cfg) + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + _emit(sd, p + "mlp.gate.weight", mlp.gate.weight, quant_cfg) + _emit(sd, p + "mlp.gate.e_score_correction_bias", + mlp.gate.e_score_correction_bias, quant_cfg) + gup, dwn = mlp.experts.gate_up_proj, mlp.experts.down_proj + for e in range(cfg.n_routed_experts): + _emit(sd, p + f"mlp.experts.{e}.gate_proj.weight", + gup[e, :moe_inter, :], quant_cfg) + _emit(sd, p + f"mlp.experts.{e}.up_proj.weight", + gup[e, moe_inter:, :], quant_cfg) + _emit(sd, p + f"mlp.experts.{e}.down_proj.weight", dwn[e], quant_cfg) + sh = mlp.shared_expert + _emit(sd, p + "mlp.shared_experts.gate_proj.weight", + sh.gate_up_proj.weight[:shared_inter], quant_cfg) + _emit(sd, p + "mlp.shared_experts.up_proj.weight", + sh.gate_up_proj.weight[shared_inter:], quant_cfg) + _emit(sd, p + "mlp.shared_experts.down_proj.weight", sh.down_proj.weight, quant_cfg) + else: + _emit(sd, p + "mlp.gate_proj.weight", mlp.gate_up_proj.weight[:inter], quant_cfg) + _emit(sd, p + "mlp.up_proj.weight", mlp.gate_up_proj.weight[inter:], quant_cfg) + _emit(sd, p + "mlp.down_proj.weight", mlp.down_proj.weight, quant_cfg) + _emit(sd, "model.norm.weight", m.norm.weight, quant_cfg) + _emit(sd, "lm_head.weight", model.lm_head.weight, quant_cfg) + return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()} + + +def _build_loaded(cfg, checkpoint_dir): + """Production path: meta -> to(bf16) -> to_empty(cuda) -> load_weights.""" + with torch.device("meta"): + model = KimiForCausalLM(cfg) + model = model.to(DTYPE) + model.to_empty(device=DEVICE) + loaded = driver_load_weights(model, checkpoint_dir, device=DEVICE) + return model.eval(), loaded + + +# -------------------------------------------------------------------------- +# Test 1 — full-model packed-expert load + forward vs dequant-on-load. +# -------------------------------------------------------------------------- + +def test_inkernel_weight_loading_and_forward_vs_dequant_on_load(tmp_path): + from safetensors.torch import save_file + + torch.manual_seed(0) + cfg_a = KimiK2Config.reduced_quantized() # dequant-on-load (bf16 experts) + cfg_b = KimiK2Config.reduced_quantized_inkernel() # in-kernel dequant (packed experts) + assert cfg_b.moe_in_kernel_dequant and cfg_b.quantization_config is not None + + ref = _build_reference(cfg_a) + assert not isinstance(ref.model.layers[0].mlp, KimiSparseMoeBlock) # dense + assert isinstance(ref.model.layers[1].mlp, KimiSparseMoeBlock) # MoE + + sd = _hf_quant_checkpoint(ref, cfg_a, cfg_a.quantization_config) + # The routed experts really are packed in the checkpoint (else in-kernel dequant is a no-op). + assert any(k.endswith("mlp.experts.0.gate_proj.weight_packed") for k in sd) + assert any(k.endswith("mlp.experts.0.down_proj.weight_packed") for k in sd) + save_file(sd, str(tmp_path / "model.safetensors")) + + model_a, _ = _build_loaded(cfg_a, tmp_path) # dequant-on-load reference + model_b, loaded_b = _build_loaded(cfg_b, tmp_path) # packed experts under test + + # (a) completeness: packed-expert loaded set == its named_parameters — includes the + # packed/scale params and EXCLUDES the bf16 fused expert params. + all_params_b = set(dict(model_b.named_parameters()).keys()) + assert loaded_b == all_params_b, ( + f"unloaded: {all_params_b - loaded_b}; spurious: {loaded_b - all_params_b}") + moe_prefix = "model.layers.1.mlp.experts." + for suffix in ("gate_up_proj_packed", "gate_up_proj_scale", + "down_proj_packed", "down_proj_scale"): + assert moe_prefix + suffix in all_params_b, f"missing {suffix}" + assert moe_prefix + "gate_up_proj" not in all_params_b # bf16 fused param gone + assert moe_prefix + "down_proj" not in all_params_b + + # (b) downcast-exemption guard: the packed params survived meta -> to(bf16) -> + # to_empty as int32 (PyTorch .to(dtype) skips integer tensors); scales are bf16. + experts_b = model_b.model.layers[1].mlp.experts + assert experts_b.gate_up_proj_packed.dtype == torch.int32 + assert experts_b.down_proj_packed.dtype == torch.int32 + assert experts_b.gate_up_proj_scale.dtype == DTYPE + assert experts_b.down_proj_scale.dtype == DTYPE + + # (c) router bias fp32; no stray buffers survived either load path. + assert model_b.model.layers[1].mlp.gate.e_score_correction_bias.dtype == torch.float32 + assert {n for n, _ in model_b.named_buffers()} == set() + + # Sanity: the two models' shared (non-expert) params are bit-identical, so any + # forward difference is isolated to the routed-expert path (in-kernel dequant vs dequant-on-load). + a_sd = dict(model_a.named_parameters()) + b_sd = dict(model_b.named_parameters()) + shared_keys = set(a_sd) & set(b_sd) + assert "model.layers.1.mlp.gate.weight" in shared_keys + for name in shared_keys: + assert torch.equal(a_sd[name], b_sd[name]), f"shared param mismatch at {name}" + + # (d) full forward: in-kernel dequant matches dequant-on-load within a loose bf16 + # tolerance (they dequant to identical bf16 weights; residual is accumulation order only). + T = 8 + ids = torch.randint(0, cfg_b.vocab_size, (T,), device=DEVICE) + pos = torch.arange(T, device=DEVICE) + with torch.no_grad(): + got = model_b(ids, _MockMLACache(cfg_b.padded_head_dim), pos) + expected = model_a(ids, _MockMLACache(cfg_a.padded_head_dim), pos) + assert got.shape == (T, cfg_b.vocab_size) + torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2) + + +# -------------------------------------------------------------------------- +# Test 2 — tp=2 packed-MoE-block simulation (packed loaders' TP slicing + +# _dispatch_packed_experts all-reduce path). Mirrors test_kimi_tp.py's block sim. +# -------------------------------------------------------------------------- + +class _NoCommGroup(CommGroup): + """world_size=2 comm group whose collectives are LOCAL no-ops: each rank + returns its partial and the test sums the two to reconstruct the reduce.""" + + def __init__(self, rank: int) -> None: + super().__init__(my_global_rank=rank, my_group_rank=rank, group_members=[0, 1]) + + def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: + return input_ + + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: + return input_ + + +def _packed_source(cfg, seed): + """Full-size expert source: bf16 gate/up/down, each fake-quantized to a full + packed+scale pair the per-rank loader then slices.""" + g = torch.Generator().manual_seed(seed) + E, H, I = cfg.n_routed_experts, cfg.hidden_size, cfg.moe_intermediate_size + sh = I * cfg.n_shared_experts + qc = cfg.quantization_config + + def rn(*shape, std=0.05, mean=0.0): + return torch.randn(*shape, generator=g) * std + mean + + def quant_stack(w): # (E, N, K) -> full packed (E,N,K//pf), scale (E,N,K//gs) + packs, scales = [], [] + for e in range(w.shape[0]): + pk, sc, _ = fake_quantize_weight( + w[e], num_bits=qc.num_bits, group_size=qc.group_size, + symmetric=qc.symmetric, scale_dtype=DTYPE) + packs.append(pk) + scales.append(sc) + return torch.stack(packs), torch.stack(scales) + + gate = rn(E, I, H, std=0.2) + up = rn(E, I, H, std=0.2) + down = rn(E, H, I, std=0.2) + return { + "router_w": torch.randn(E, H, generator=g), + "router_b": torch.randn(E, generator=g), + "gate_packed": quant_stack(gate), + "up_packed": quant_stack(up), + "down_packed": quant_stack(down), + "sh_gate": rn(sh, H), "sh_up": rn(sh, H), "sh_down": rn(H, sh), + } + + +def _load_moe_packed(block, src): + block.gate.weight.data = src["router_w"].to(DEVICE) + block.gate.e_score_correction_bias.data = src["router_b"].to(DEVICE) + gup_p, gup_s = block.experts.gate_up_proj_packed, block.experts.gate_up_proj_scale + dwn_p, dwn_s = block.experts.down_proj_packed, block.experts.down_proj_scale + gate_pk, gate_sc = src["gate_packed"] + up_pk, up_sc = src["up_packed"] + down_pk, down_sc = src["down_packed"] + for e in range(block.num_experts): + gup_p.weight_loader(gup_p, gate_pk[e].to(DEVICE), loaded_shard_id=f"gate:{e}") + gup_p.weight_loader(gup_p, up_pk[e].to(DEVICE), loaded_shard_id=f"up:{e}") + gup_s.weight_loader(gup_s, gate_sc[e].to(DEVICE), loaded_shard_id=f"gate:{e}") + gup_s.weight_loader(gup_s, up_sc[e].to(DEVICE), loaded_shard_id=f"up:{e}") + dwn_p.weight_loader(dwn_p, down_pk[e].to(DEVICE), loaded_shard_id=f"down:{e}") + dwn_s.weight_loader(dwn_s, down_sc[e].to(DEVICE), loaded_shard_id=f"down:{e}") + s = block.shared_expert + s.gate_up_proj.weight.weight_loader( + s.gate_up_proj.weight, src["sh_gate"].to(DEVICE, DTYPE), loaded_shard_id=0) + s.gate_up_proj.weight.weight_loader( + s.gate_up_proj.weight, src["sh_up"].to(DEVICE, DTYPE), loaded_shard_id=1) + s.down_proj.weight.weight_loader(s.down_proj.weight, src["sh_down"].to(DEVICE, DTYPE)) + + +def test_packed_moe_block_tp2_sim_matches_tp1(): + cfg = KimiK2Config.reduced_quantized_inkernel() # shard_inter=32 at tp=2 (ok) + src = _packed_source(cfg, seed=707) + g = torch.Generator().manual_seed(808) + h = (torch.randn(7, cfg.hidden_size, generator=g) * 0.1).to(DEVICE, DTYPE) + + ref = KimiSparseMoeBlock(cfg, CommGroup.trivial()).to(DEVICE, DTYPE) + _load_moe_packed(ref, src) + full_inter = ref.experts.gate_up_proj_packed.shape[1] # 2*moe_inter + out_ref = ref(h) + + partials = [] + for rank in range(2): + block = KimiSparseMoeBlock(cfg, _NoCommGroup(rank)).to(DEVICE, DTYPE) + # each rank holds only a 1/2 stripe of the fused (packed) intermediate + assert block.experts.gate_up_proj_packed.shape[1] == full_inter // 2 + assert block.experts.down_proj_packed.dtype == torch.int32 + _load_moe_packed(block, src) + partials.append(block(h)) + + out_tp2 = partials[0] + partials[1] # intermediate-parallel reduce == sum of ranks + max_abs = (out_tp2 - out_ref).abs().max().item() + assert max_abs < 5e-2, f"packed MoE tp2 vs tp1 max abs diff {max_abs}" + torch.testing.assert_close(out_tp2, out_ref, rtol=2e-2, atol=2e-2) diff --git a/test/integration/test_kimi_quant_weight_loading.py b/test/integration/test_kimi_quant_weight_loading.py new file mode 100644 index 000000000..e6b19db32 --- /dev/null +++ b/test/integration/test_kimi_quant_weight_loading.py @@ -0,0 +1,269 @@ +"""Golden: compressed-tensors dequant-on-load for Kimi-K2.7. + +The real 1T INT4 checkpoint is absent and a bf16 dequant of it would not fit, so +this validates the *parser + numerics* on a SYNTHETIC ``reduced_quantized()`` +model, exactly mirroring the bf16 ``test_kimi_weight_loading.py`` but with a +compressed-tensors quantized checkpoint: + + 1. build a ``KimiForCausalLM`` reference with random bf16 weights; + 2. **fake-quantize in place** every eligible weight (2-D, input dim divisible by + ``group_size``, and not a norm / router / embedding / lm_head) — writing the + dequantized bf16 back into the reference, so the reference now holds exactly + what a correct loader must reproduce — and serialize it as an HF + compressed-tensors checkpoint: ``.weight_packed`` (int32) + + ``.weight_scale`` (bf16) for quantized weights, plain ``.weight`` + for the rest (the ``ignore`` set + weights whose dim doesn't divide the + group). MLA / dense-FFN / shared-expert / routed-expert weights are all + quantized here — proving dequant-on-load covers plain linears *and* the fused experts; + 3. build a fresh model on ``meta`` -> ``to(bf16)`` -> ``to_empty(cuda)`` and load + the checkpoint via the standard driver. Because the config carries a + ``quantization_config``, ``load_kimi_hf_weights`` wraps the stream with the + dequant-on-load dequantizer *before* the remap + stacked rules (which are unchanged); + 4. assert (a) completeness, (b) every loaded param equals the fake-quantized + reference bit-for-bit (the loader reproduces the dequant exactly), (c) the + router bias stays fp32 and no stray buffers survive, and (d) a full forward on + the loaded model matches a forward on the reference model. + +Run: pytest test/integration/test_kimi_quant_weight_loading.py -v +""" +import pytest +import torch + +from mstar.model.kimi_k2_7._testing import fake_quantize_weight +from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM +from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock +from mstar.model.kimi_k2_7.config import KimiK2Config +from mstar.model.loader import load_weights as driver_load_weights + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Kimi quant weight-loading golden needs a GPU (RMSNorm + fused expert GEMM)", +) + +DEVICE = "cuda" + + +# -------------------------------------------------------------------------- +# Mock paged cache (causal SDPA at 1/sqrt(head_dim)) — same as test_kimi_forward. +# -------------------------------------------------------------------------- + +def _sdpa_causal(q, k, v, scale): + qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) + T = q.shape[0] + causal = torch.triu( + torch.full((T, T), float("-inf"), device=q.device), diagonal=1) + attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1) + return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) + + +class _MockMLACache: + def __init__(self, head_dim): + self.scale = head_dim ** -0.5 + + def set_layer_idx(self, _i): + pass + + def advance_seq_lens(self, *_a, **_k): + pass + + def run_attention(self, q, k, v): + return _sdpa_causal(q, k, v, self.scale) + + +# -------------------------------------------------------------------------- +# Random weight init (router bias kept fp32) — same as test_kimi_weight_loading. +# -------------------------------------------------------------------------- + +def _fill_layer(layer, cfg): + a = layer.self_attn + for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): + lin.weight.data.normal_(0, 0.03) + for norm in (a.q_a_layernorm, a.kv_a_layernorm): + norm.weight.data.normal_(1.0, 0.02) + layer.input_layernorm.weight.data.normal_(1.0, 0.02) + layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02) + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + mlp.gate.weight.data.normal_(0, 1) + mlp.gate.e_score_correction_bias.data = torch.randn( + cfg.n_routed_experts, device=DEVICE, dtype=torch.float32) + mlp.experts.gate_up_proj.data.normal_(0, 0.05) + mlp.experts.down_proj.data.normal_(0, 0.05) + mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05) + else: + mlp.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.down_proj.weight.data.normal_(0, 0.05) + + +def _build_reference(cfg): + model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=torch.bfloat16) + model.model.embed_tokens.weight.data.normal_(0, 0.05) + model.model.norm.weight.data.normal_(1.0, 0.02) + model.lm_head.weight.data.normal_(0, 0.02) + for layer in model.model.layers: + _fill_layer(layer, cfg) + # Eval-only fixture; disable grad so the in-place dequant write-back in _emit + # (copying into leaf params) is allowed. + model.requires_grad_(False) + return model.eval() + + +# -------------------------------------------------------------------------- +# Fake-quant serialization: compressed-tensors keys for eligible weights, plain +# bf16 for the rest. Mutates the reference in place to hold the dequant. +# -------------------------------------------------------------------------- + +def _keep_bf16(key): + """Weights a real compressed-tensors Kimi checkpoint leaves in bf16.""" + if key == "lm_head.weight": + return True + return ( + key.endswith("norm.weight") # all RMSNorms incl. *_layernorm + or key.endswith("mlp.gate.weight") # MoE router — never quantized + or key.endswith("e_score_correction_bias") + or key.endswith("embed_tokens.weight") + ) + + +def _emit(sd, key, view, quant_cfg): + """Add checkpoint entries for reference param ``view`` at ``key``. + + Quantizes (and writes the dequant back into ``view``) when eligible, else + stores the plain bf16 tensor. Eligibility mirrors a real checkpoint: 2-D, + input dim divisible by the group size, and not in the bf16-keep set. + """ + eligible = ( + not _keep_bf16(key) + and view.dim() == 2 + and view.shape[-1] % quant_cfg.group_size == 0 + ) + if not eligible: + sd[key] = view + return + # Store the scale in bf16 (as a real compressed-tensors checkpoint does) and + # take the dequant from that same bf16 scale, so the reference holds exactly + # what the loader's bf16-scale dequant reconstructs. + packed, scale, deq = fake_quantize_weight( + view, num_bits=quant_cfg.num_bits, group_size=quant_cfg.group_size, + symmetric=quant_cfg.symmetric, scale_dtype=torch.bfloat16, + ) + base = key[: -len(".weight")] + sd[base + ".weight_packed"] = packed + sd[base + ".weight_scale"] = scale + view.copy_(deq.to(view.dtype)) # reference now holds the dequantized weight + + +def _hf_quant_checkpoint(model, cfg, quant_cfg): + inter = cfg.intermediate_size + moe_inter = cfg.moe_intermediate_size + shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts + m = model.model + sd = {} + _emit(sd, "model.embed_tokens.weight", m.embed_tokens.weight, quant_cfg) + for i, layer in enumerate(m.layers): + p = f"model.layers.{i}." + a = layer.self_attn + _emit(sd, p + "self_attn.q_a_proj.weight", a.q_a_proj.weight, quant_cfg) + _emit(sd, p + "self_attn.q_a_layernorm.weight", a.q_a_layernorm.weight, quant_cfg) + _emit(sd, p + "self_attn.q_b_proj.weight", a.q_b_proj.weight, quant_cfg) + _emit(sd, p + "self_attn.kv_a_proj_with_mqa.weight", + a.kv_a_proj_with_mqa.weight, quant_cfg) + _emit(sd, p + "self_attn.kv_a_layernorm.weight", a.kv_a_layernorm.weight, quant_cfg) + _emit(sd, p + "self_attn.kv_b_proj.weight", a.kv_b_proj.weight, quant_cfg) + _emit(sd, p + "self_attn.o_proj.weight", a.o_proj.weight, quant_cfg) + _emit(sd, p + "input_layernorm.weight", layer.input_layernorm.weight, quant_cfg) + _emit(sd, p + "post_attention_layernorm.weight", + layer.post_attention_layernorm.weight, quant_cfg) + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + _emit(sd, p + "mlp.gate.weight", mlp.gate.weight, quant_cfg) + _emit(sd, p + "mlp.gate.e_score_correction_bias", + mlp.gate.e_score_correction_bias, quant_cfg) + gup, dwn = mlp.experts.gate_up_proj, mlp.experts.down_proj + for e in range(cfg.n_routed_experts): + _emit(sd, p + f"mlp.experts.{e}.gate_proj.weight", + gup[e, :moe_inter, :], quant_cfg) + _emit(sd, p + f"mlp.experts.{e}.up_proj.weight", + gup[e, moe_inter:, :], quant_cfg) + _emit(sd, p + f"mlp.experts.{e}.down_proj.weight", dwn[e], quant_cfg) + sh = mlp.shared_expert + _emit(sd, p + "mlp.shared_experts.gate_proj.weight", + sh.gate_up_proj.weight[:shared_inter], quant_cfg) + _emit(sd, p + "mlp.shared_experts.up_proj.weight", + sh.gate_up_proj.weight[shared_inter:], quant_cfg) + _emit(sd, p + "mlp.shared_experts.down_proj.weight", sh.down_proj.weight, quant_cfg) + else: + _emit(sd, p + "mlp.gate_proj.weight", mlp.gate_up_proj.weight[:inter], quant_cfg) + _emit(sd, p + "mlp.up_proj.weight", mlp.gate_up_proj.weight[inter:], quant_cfg) + _emit(sd, p + "mlp.down_proj.weight", mlp.down_proj.weight, quant_cfg) + _emit(sd, "model.norm.weight", m.norm.weight, quant_cfg) + _emit(sd, "lm_head.weight", model.lm_head.weight, quant_cfg) + # Clone to cpu + break storage aliasing (safetensors rejects shared storage). + return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()} + + +def _build_loaded(cfg, checkpoint_dir): + """Production path: meta -> to(bf16) -> to_empty(cuda) -> load_weights.""" + with torch.device("meta"): + model = KimiForCausalLM(cfg) + model = model.to(torch.bfloat16) + model.to_empty(device=DEVICE) + loaded = driver_load_weights(model, checkpoint_dir, device=DEVICE) + return model.eval(), loaded + + +# -------------------------------------------------------------------------- +# Test +# -------------------------------------------------------------------------- + +def test_quant_weight_loading_roundtrip_and_forward(tmp_path): + from safetensors.torch import save_file + + torch.manual_seed(0) + cfg = KimiK2Config.reduced_quantized() # group_size=32, INT4 symmetric + assert cfg.quantization_config is not None + ref = _build_reference(cfg) + # The stack spans the dense->MoE transition (first_k_dense_replace=1). + assert not isinstance(ref.model.layers[0].mlp, KimiSparseMoeBlock) + assert isinstance(ref.model.layers[1].mlp, KimiSparseMoeBlock) + + sd = _hf_quant_checkpoint(ref, cfg, cfg.quantization_config) + + # Guard: the checkpoint really is quantized (else this silently degrades to + # the bf16 test) — routed experts AND plain linears carry packed weights. + assert any(k.endswith("mlp.experts.0.gate_proj.weight_packed") for k in sd) + assert any(k.endswith("self_attn.o_proj.weight_packed") for k in sd) + # ... and the ignore-set weights stayed bf16 (plain .weight, no packed). + assert "lm_head.weight" in sd and "lm_head.weight_packed" not in sd + assert "model.embed_tokens.weight" in sd + + save_file(sd, str(tmp_path / "model.safetensors")) + model, loaded = _build_loaded(cfg, tmp_path) + + # (a) completeness: every param received exactly one tensor, and no quant + # sub-key leaked through as a spurious param. + all_params = set(dict(model.named_parameters()).keys()) + assert loaded == all_params, ( + f"unloaded: {all_params - loaded}; spurious: {loaded - all_params}") + + # (b) every loaded param equals the fake-quantized reference, bit for bit + # (dequant-on-load reproduces the same fp32 (q-bias)*scale -> bf16 dequant). + ref_sd = dict(ref.named_parameters()) + for name, param in model.named_parameters(): + assert torch.equal(param, ref_sd[name]), f"mismatch at {name}" + + # (c) router bias preserved fp32; no derived buffers survived the load path. + bias = model.model.layers[1].mlp.gate.e_score_correction_bias + assert bias.dtype == torch.float32 + assert {n for n, _ in model.named_buffers()} == set() + + # (d) full forward on the loaded model matches the reference model's forward. + T = 8 + ids = torch.randint(0, cfg.vocab_size, (T,), device=DEVICE) + pos = torch.arange(T, device=DEVICE) + with torch.no_grad(): + got = model(ids, _MockMLACache(cfg.padded_head_dim), pos) + expected = ref(ids, _MockMLACache(cfg.padded_head_dim), pos) + assert got.shape == (T, cfg.vocab_size) + torch.testing.assert_close(got, expected, rtol=1e-3, atol=1e-3) diff --git a/test/modular/test_kimi_k27_code_wiring.py b/test/modular/test_kimi_k27_code_wiring.py new file mode 100644 index 000000000..90c270b75 --- /dev/null +++ b/test/modular/test_kimi_k27_code_wiring.py @@ -0,0 +1,220 @@ +"""CPU wiring tests for the REAL ``moonshotai/Kimi-K2.7-Code`` text-only serve. + +No weights, no GPU — the golden gate for the K2.7-Code serve plumbing while the +595 GB checkpoint is still downloading. Three concerns: + + 1. :meth:`KimiK2Config.k27_code` builds the full-size 1T text config with packed + experts armed and — crucially — keeps the default ``beta_fast=32.0`` (the + K2.7-Code ``text_config`` value). Guards against clobbering the YaRN field. + 2. :func:`kimi_name_remapper` strips the multimodal ``language_model.`` prefix so + the DeepSeek-V3 text keys land on ``KimiForCausalLM``'s params, routes the + packed routed-expert sub-keys through the packed-expert stacked rules, and drops the + vision (``vision_tower.*`` / ``mm_projector.*``) + ``weight_shape`` keys. A bare + ``model.*`` key (no prefix) is left unchanged. + 3. ``_maybe_apply_checkpoint_quant_config`` reads the ``quantization_config`` + NESTED under ``text_config`` (the multimodal wrapper leaves the top-level + null), while still parsing a flat top-level block and staying ``None`` for a + plain-bf16 checkpoint. + +Run: pytest test/modular/test_kimi_k27_code_wiring.py -v +""" +import json + +import torch + +from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM +from mstar.model.kimi_k2_7.config import KimiK2Config +from mstar.model.kimi_k2_7.kimi_model import KimiK2Model +from mstar.model.kimi_k2_7.weight_loader import ( + build_kimi_stacked_params, + kimi_name_remapper, +) +from mstar.model.loader.base import _apply_stacked + +# -------------------------------------------------------------------------- +# 1. Config: k27_code() == full 1T dims + packed experts, default beta_fast=32.0. +# -------------------------------------------------------------------------- + +def test_k27_code_config_full_dims_packed_and_beta_fast(): + cfg = KimiK2Config.k27_code() + + # Packed experts armed, quant config auto-read from the checkpoint (still None here). + assert cfg.moe_in_kernel_dequant is True + assert cfg.quantization_config is None + + # K2.7-Code keeps beta_fast=32.0 (guard against clobbering the YaRN field). + assert cfg.rope_scaling["beta_fast"] == 32.0 + assert cfg.rope_scaling["factor"] == 64.0 + assert cfg.rope_scaling["rope_type"] == "deepseek_yarn" + + # Full 1T text dims, matching the real Kimi-K2.7-Code text_config. + assert cfg.num_hidden_layers == 61 + assert cfg.n_routed_experts == 384 + assert cfg.hidden_size == 7168 + assert cfg.q_lora_rank == 1536 + assert cfg.kv_lora_rank == 512 + assert cfg.moe_intermediate_size == 2048 + assert cfg.routed_scaling_factor == 2.827 + assert cfg.qk_nope_head_dim == 128 + assert cfg.qk_rope_head_dim == 64 + assert cfg.v_head_dim == 128 + + # It really is the full-size default plus exactly the one flag (no dim drift). + base = KimiK2Config() + assert cfg.num_hidden_layers == base.num_hidden_layers + assert cfg.rope_scaling == base.rope_scaling # NO beta_fast override + assert base.moe_in_kernel_dequant is False and cfg.moe_in_kernel_dequant is True + + +# -------------------------------------------------------------------------- +# 2. Remapper: language_model.* strip + packed-expert routing + drops. +# -------------------------------------------------------------------------- + +def _route(name, stacked): + """Mirror the loader: name_remapper then stacked-shard routing.""" + mapped = kimi_name_remapper(name) + if mapped is None: + return None, None + return _apply_stacked(mapped, stacked) + + +def test_remapper_language_model_prefix_and_packed_experts(): + cfg = KimiK2Config.reduced_quantized_inkernel() # in-kernel dequant => packed expert params + with torch.device("meta"): + model = KimiForCausalLM(cfg) + params = set(dict(model.named_parameters()).keys()) + stacked = build_kimi_stacked_params(cfg.n_routed_experts, packed_experts=True) + + # -- plain text keys: strip language_model., land on a real param ------------ + assert kimi_name_remapper( + "language_model.model.layers.0.self_attn.q_a_proj.weight" + ) == "model.layers.0.self_attn.q_a_proj.weight" + assert ( + kimi_name_remapper("language_model.model.embed_tokens.weight") + == "model.embed_tokens.weight" + ) + assert kimi_name_remapper("language_model.lm_head.weight") == "lm_head.weight" + for landed in ( + "model.layers.0.self_attn.q_a_proj.weight", + "model.embed_tokens.weight", + "lm_head.weight", + ): + assert landed in params + + # -- shared expert: plural -> singular --------------------------------------- + shared = kimi_name_remapper( + "language_model.model.layers.1.mlp.shared_experts.down_proj.weight" + ) + assert shared == "model.layers.1.mlp.shared_expert.down_proj.weight" + assert shared in params + + # -- packed routed expert: remap + stacked -> the FOUR packed params --------- + gate_p, gate_sid = _route( + "language_model.model.layers.1.mlp.experts.0.gate_proj.weight_packed", stacked + ) + assert gate_p == "model.layers.1.mlp.experts.gate_up_proj_packed" + assert gate_sid == "gate:0" + assert gate_p in params + + scale_p, scale_sid = _route( + "language_model.model.layers.1.mlp.experts.0.gate_proj.weight_scale", stacked + ) + assert scale_p == "model.layers.1.mlp.experts.gate_up_proj_scale" + assert scale_sid == "gate:0" + assert scale_p in params + + down_p, down_sid = _route( + "language_model.model.layers.1.mlp.experts.0.down_proj.weight_packed", stacked + ) + assert down_p == "model.layers.1.mlp.experts.down_proj_packed" + assert down_sid == "down:0" + assert down_p in params + + # -- vision drop: identity remap, NOT a model param (base loader skips it) ---- + for vkey in ( + "vision_tower.encoder.blocks.0.wqkv.weight", + "mm_projector.proj.0.weight", + ): + assert kimi_name_remapper(vkey) == vkey # identity — no surgery + target, _ = _route(vkey, stacked) + assert target not in params # dropped + + # -- weight_shape drop: routes to no real param ------------------------------ + ws_target, _ = _route( + "language_model.model.layers.1.mlp.experts.0.gate_proj.weight_shape", stacked + ) + assert ws_target not in params + + # -- a flat model.* key (no language_model. prefix) is untouched ------------- + assert ( + kimi_name_remapper("model.layers.0.self_attn.q_a_proj.weight") + == "model.layers.0.self_attn.q_a_proj.weight" + ) + + +# -------------------------------------------------------------------------- +# 3. Nested-quant reader: text_config.quantization_config + flat + bf16. +# -------------------------------------------------------------------------- + +_QUANT_BLOCK = { + "format": "pack-quantized", + "quant_method": "compressed-tensors", + "ignore": ["lm_head", "re:.*self_attn.*", "re:.*shared_experts.*"], + "config_groups": { + "group_0": { + "weights": { + "num_bits": 4, + "group_size": 32, + "symmetric": True, + "strategy": "group", + "type": "int", + }, + "targets": ["Linear"], + } + }, +} + + +def _make_model_with_config_json(tmp_dir, config_dict): + """A KimiK2Model with a bf16 (quant=None) config and a written config.json.""" + (tmp_dir / "config.json").write_text(json.dumps(config_dict)) + model = object.__new__(KimiK2Model) + model.config = KimiK2Config() # quantization_config is None by default + return model + + +def test_nested_quant_config_read(tmp_path): + # Nested under text_config, top-level absent — the real K2.7-Code layout. + d = tmp_path / "nested" + d.mkdir() + model = _make_model_with_config_json( + d, {"text_config": {"num_hidden_layers": 61, "quantization_config": _QUANT_BLOCK}} + ) + model._maybe_apply_checkpoint_quant_config(str(d)) + qc = model.config.quantization_config + assert qc is not None + assert qc.num_bits == 4 + assert qc.group_size == 32 + assert qc.symmetric is True + assert qc.quant_format == "pack-quantized" + + +def test_flat_quant_config_read_backward_compat(tmp_path): + # A flat top-level quantization_config block must still parse. + d = tmp_path / "flat" + d.mkdir() + model = _make_model_with_config_json(d, {"quantization_config": _QUANT_BLOCK}) + model._maybe_apply_checkpoint_quant_config(str(d)) + qc = model.config.quantization_config + assert qc is not None + assert qc.num_bits == 4 + assert qc.group_size == 32 + + +def test_plain_bf16_config_stays_none(tmp_path): + # No quant block anywhere (nested or flat) — stays bf16 (None). + d = tmp_path / "bf16" + d.mkdir() + model = _make_model_with_config_json(d, {"text_config": {"num_hidden_layers": 61}}) + model._maybe_apply_checkpoint_quant_config(str(d)) + assert model.config.quantization_config is None diff --git a/test/modular/test_kimi_quant.py b/test/modular/test_kimi_quant.py new file mode 100644 index 000000000..aed5f4df3 --- /dev/null +++ b/test/modular/test_kimi_quant.py @@ -0,0 +1,278 @@ +"""CPU unit tests for the Kimi-K2.7 compressed-tensors dequant utilities (dequant-on-load). + +These pin the *numerics and bit layout* of the dequant-on-load parser without a +GPU or checkpoint — the cheapest level that guards the correctness harness: + + 1. known-answer pack/unpack (the exact int32 bit layout, incl. the sign-bit + nibble), so a future refactor can't silently change the on-disk convention; + 2. pack/unpack round-trip on random nibbles; + 3. symmetric dequant math (offset-binary ``(nibble - bias) * scale``); + 4. ``fake_quantize_weight`` -> ``dequantize_weight`` exactness (the golden + harness relies on the loader reproducing the fake-quant result bit-for-bit); + 5. the streaming generator: quant components collapse to one bf16 ``*.weight``, + non-quant keys pass through, incomplete groups raise; + 6. ``CompressedTensorsQuantConfig.from_hf_config_dict`` parsing. + +The full weight-loading + forward golden (needs the fused-expert GEMM + RMSNorm) +lives in ``test/integration/test_kimi_quant_weight_loading.py``. + +Run: pytest test/modular/test_kimi_quant.py -v +""" +import pytest +import torch + +from mstar.model.kimi_k2_7._testing import fake_quantize_weight +from mstar.model.kimi_k2_7.quantization import ( + CompressedTensorsQuantConfig, + dequant_compressed_tensors_stream, + dequantize_weight, + pack_int32, + unpack_int32, +) + +# -------------------------------------------------------------------------- +# 1. Known-answer pack/unpack — pins the int32 bit layout. +# -------------------------------------------------------------------------- + +def test_pack_known_answer(): + # Eight INT4 nibbles 0..7 along the last axis pack low-order-first: + # sum(j << 4*j for j in 0..7) == 0x76543210. + nibbles = torch.arange(8, dtype=torch.int64).reshape(1, 8) + packed = pack_int32(nibbles, num_bits=4) + assert packed.dtype == torch.int32 + assert packed.shape == (1, 1) + assert packed.item() == 0x76543210 + + # A top nibble >= 8 sets bit 31, so the int32 container is negative — the + # unpack must still recover it (reads the 32-bit pattern as unsigned). + top = torch.tensor([[0, 0, 0, 0, 0, 0, 0, 8]], dtype=torch.int64) + packed_top = pack_int32(top, num_bits=4) + assert packed_top.item() == -(2**31) # 0x80000000 as signed int32 + back = unpack_int32(packed_top, num_bits=4) + assert torch.equal(back, top) + + +def test_pack_unpack_roundtrip(): + torch.manual_seed(0) + nibbles = torch.randint(0, 16, (5, 32), dtype=torch.int64) # in=32 -> packed 4 + packed = pack_int32(nibbles, num_bits=4) + assert packed.shape == (5, 4) + assert torch.equal(unpack_int32(packed, num_bits=4), nibbles) + + +# -------------------------------------------------------------------------- +# 2. Dequant math — symmetric offset-binary (nibble - bias) * scale. +# -------------------------------------------------------------------------- + +def test_dequantize_symmetric_known_answer(): + # One row, one group of 8 (group_size=8). Unsigned nibbles minus bias 8 give + # the signed quantized values, times a per-group scale of 2.0. + unsigned = torch.tensor([[8, 9, 7, 8, 10, 6, 8, 8]], dtype=torch.int64) + signed = torch.tensor([[0, 1, -1, 0, 2, -2, 0, 0]], dtype=torch.float32) + packed = pack_int32(unsigned, num_bits=4) + scale = torch.tensor([[2.0]]) # (out=1, groups=1) + got = dequantize_weight( + packed, scale, num_bits=4, group_size=8, symmetric=True, + out_dtype=torch.float32, + ) + assert torch.equal(got, signed * 2.0) + + +def test_dequantize_two_groups_broadcast(): + # in=16, group_size=8 -> two groups with distinct scales; check the scale + # broadcasts per-group along the input axis. + unsigned = torch.full((1, 16), 8, dtype=torch.int64) + unsigned[0, 0] = 9 # +1 in group 0 + unsigned[0, 8] = 9 # +1 in group 1 + packed = pack_int32(unsigned, num_bits=4) + scale = torch.tensor([[3.0, 5.0]]) # group0=3, group1=5 + got = dequantize_weight( + packed, scale, num_bits=4, group_size=8, symmetric=True, + out_dtype=torch.float32, + ) + assert got[0, 0] == 3.0 + assert got[0, 8] == 5.0 + assert got[0, 1] == 0.0 + + +# -------------------------------------------------------------------------- +# 3. fake_quantize -> dequantize exactness (the golden's core invariant). +# -------------------------------------------------------------------------- + +@pytest.mark.parametrize("group_size", [8, 16, -1]) +def test_fake_quantize_dequantize_exact(group_size): + torch.manual_seed(1) + w = torch.randn(12, 32) * 0.1 + packed, scale, deq = fake_quantize_weight( + w, num_bits=4, group_size=group_size, symmetric=True, + ) + # Reconstructing from the on-disk tensors must reproduce the fake-quant bf16 + # result bit-for-bit (both compute q*scale in fp32 then cast to bf16). + got = dequantize_weight( + packed, scale, num_bits=4, group_size=group_size, symmetric=True, + ) + assert got.dtype == torch.bfloat16 + assert torch.equal(got, deq) + # And the quantization is lossy but bounded (sanity: close to the original). + assert (got.float() - w).abs().max() < 0.05 + + +# -------------------------------------------------------------------------- +# 4. The dequant-on-load streaming generator. +# -------------------------------------------------------------------------- + +def _quant_components(base, w, cfg): + # Store the scale in bf16 (as a real compressed-tensors checkpoint does); the + # returned dequant is derived from that same bf16 scale, so it matches the + # stream's reconstruction bit-for-bit. + packed, scale, deq = fake_quantize_weight( + w, num_bits=cfg.num_bits, group_size=cfg.group_size, + symmetric=cfg.symmetric, scale_dtype=torch.bfloat16, + ) + return { + f"{base}.weight_packed": packed, + f"{base}.weight_scale": scale, + f"{base}.weight_shape": torch.tensor(list(w.shape), dtype=torch.int64), + }, deq + + +def test_stream_dequantizes_and_passes_through(): + torch.manual_seed(2) + cfg = CompressedTensorsQuantConfig(num_bits=4, group_size=16, symmetric=True) + w_a = torch.randn(8, 32) * 0.1 + w_b = torch.randn(4, 16) * 0.1 + comp_a, deq_a = _quant_components("layer.0.a_proj", w_a, cfg) + comp_b, deq_b = _quant_components("layer.0.b_proj", w_b, cfg) + norm = torch.randn(8) # a non-quant key that must pass straight through + + # Interleave/shuffle keys — the generator must reassemble by base name, + # independent of order. + stream = [ + ("layer.0.a_proj.weight_scale", comp_a["layer.0.a_proj.weight_scale"]), + ("layer.0.norm.weight", norm), + ("layer.0.b_proj.weight_packed", comp_b["layer.0.b_proj.weight_packed"]), + ("layer.0.a_proj.weight_shape", comp_a["layer.0.a_proj.weight_shape"]), + ("layer.0.a_proj.weight_packed", comp_a["layer.0.a_proj.weight_packed"]), + ("layer.0.b_proj.weight_scale", comp_b["layer.0.b_proj.weight_scale"]), + ] + out = dict(dequant_compressed_tensors_stream(iter(stream), cfg)) + + # Exactly: two dequantized *.weight keys + the passthrough norm; no quant subkeys. + assert set(out) == {"layer.0.a_proj.weight", "layer.0.b_proj.weight", "layer.0.norm.weight"} + assert torch.equal(out["layer.0.a_proj.weight"], deq_a) + assert torch.equal(out["layer.0.b_proj.weight"], deq_b) + assert torch.equal(out["layer.0.norm.weight"], norm) + + +def test_stream_incomplete_group_raises(): + cfg = CompressedTensorsQuantConfig(num_bits=4, group_size=16, symmetric=True) + w = torch.randn(4, 16) * 0.1 + comp, _ = _quant_components("x.proj", w, cfg) + # Only the packed tensor, no scale -> the group can never complete. + stream = [("x.proj.weight_packed", comp["x.proj.weight_packed"])] + with pytest.raises(ValueError, match="incomplete"): + list(dequant_compressed_tensors_stream(iter(stream), cfg)) + + +# -------------------------------------------------------------------------- +# 4b. dequant-on-load + packed-expert coexistence: keep_packed passes routed experts +# through raw while MLA/dense keys still dequantize (the streaming half of the +# mixed-load path; the GPU golden proves the packed params then load + run). +# -------------------------------------------------------------------------- + +def test_stream_keep_packed_passthrough(): + cfg = CompressedTensorsQuantConfig(num_bits=4, group_size=16, symmetric=True) + exp_base = "model.layers.1.mlp.experts.3.gate_proj" # a routed expert -> packed experts + mla_base = "model.layers.1.self_attn.o_proj" # an MLA weight -> dequant-on-load + comp_exp, _ = _quant_components(exp_base, torch.randn(8, 32) * 0.1, cfg) + comp_mla, deq_mla = _quant_components(mla_base, torch.randn(4, 16) * 0.1, cfg) + + def keep_packed(base): + return ".experts.3.gate_proj" in base + + # Expert carries all three sub-keys (packed/scale/shape); MLA carries the two + # that complete a dequant (no shape) so no dangling buffer trips the end check. + stream = [ + (f"{exp_base}.weight_packed", comp_exp[f"{exp_base}.weight_packed"]), + (f"{exp_base}.weight_scale", comp_exp[f"{exp_base}.weight_scale"]), + (f"{exp_base}.weight_shape", comp_exp[f"{exp_base}.weight_shape"]), + (f"{mla_base}.weight_packed", comp_mla[f"{mla_base}.weight_packed"]), + (f"{mla_base}.weight_scale", comp_mla[f"{mla_base}.weight_scale"]), + ] + out = dict(dequant_compressed_tensors_stream(iter(stream), cfg, keep_packed=keep_packed)) + + # Routed-expert sub-keys pass through RAW — packed int32 + scale + shape, and + # crucially NO collapsed ``.weight`` (they load into the packed params instead). + assert out[f"{exp_base}.weight_packed"].dtype == torch.int32 + assert f"{exp_base}.weight_scale" in out + assert f"{exp_base}.weight_shape" in out + assert f"{exp_base}.weight" not in out + # The MLA weight still collapses to one dequantized bf16 ``.weight`` (dequant-on-load). + assert torch.equal(out[f"{mla_base}.weight"], deq_mla) + assert f"{mla_base}.weight_packed" not in out + + +# -------------------------------------------------------------------------- +# 4c. Pure-torch reference for the W4A16 kernel math (no GPU): the per-group +# ``(unpack - 8) * scale`` -> bf16 grouped GEMM the Triton kernel replicates, +# pinning the packed-K layout and the top-nibble (bit-31) sign case. +# -------------------------------------------------------------------------- + +def test_grouped_gemm_reference_math_and_top_nibble(): + torch.manual_seed(7) + N, K, gs = 6, 32, 16 # two groups along the packed K axis + # Wide init so per-group amax uses the full nibble range -> some top nibbles + # land >= 8 (int32 container bit 31 set), exercising the sign path. + w = torch.randn(N, K) * 0.5 + packed, scale, deq = fake_quantize_weight( + w, num_bits=4, group_size=gs, symmetric=True, scale_dtype=torch.bfloat16, + ) + assert packed.shape == (N, K // 8) # packed along the last (input/K) axis + assert (packed < 0).any(), "no negative container — top-nibble sign path untested" + + # Reproduce the kernel's in-register arithmetic in pure torch: unpack the + # nibble, offset-binary subtract 8, scale per group (broadcast along K). + nibbles = unpack_int32(packed, num_bits=4).to(torch.float32) # (N, K) unsigned + scale_bc = scale.to(torch.float32).repeat_interleave(gs, dim=-1) # (N, K) + manual_deq = ((nibbles - 8.0) * scale_bc).to(torch.bfloat16) + assert torch.equal(manual_deq, deq) # kernel math == dequantize_weight + + # Grouped GEMM equivalence: the packed path must produce the same y as feeding + # the bf16 dequant directly (both contract over the same bf16 weight values). + x = torch.randn(4, K) + y_manual = torch.einsum("tk,nk->tn", x, manual_deq.float()) + y_deq = torch.einsum("tk,nk->tn", x, deq.float()) + assert torch.equal(y_manual, y_deq) + + +# -------------------------------------------------------------------------- +# 5. Config parsing. +# -------------------------------------------------------------------------- + +def test_quant_config_from_hf_dict(): + raw = { + "format": "pack-quantized", + "quant_method": "compressed-tensors", + "ignore": ["lm_head", "re:.*gate$"], + "config_groups": { + "group_0": { + "weights": { + "num_bits": 4, + "group_size": 32, + "symmetric": True, + "strategy": "group", + "type": "int", + }, + "targets": ["Linear"], + } + }, + } + cfg = CompressedTensorsQuantConfig.from_hf_config_dict(raw) + assert cfg is not None + assert cfg.num_bits == 4 + assert cfg.group_size == 32 + assert cfg.symmetric is True + assert cfg.pack_factor == 8 + assert cfg.ignore == ("lm_head", "re:.*gate$") + assert CompressedTensorsQuantConfig.from_hf_config_dict(None) is None + assert CompressedTensorsQuantConfig.from_hf_config_dict({}) is None From bda6dec0000ea539717f78a13028bdd2a4d625f4 Mon Sep 17 00:00:00 2001 From: Garv Ghai Date: Sun, 26 Jul 2026 14:24:48 +0000 Subject: [PATCH 4/9] Kimi-K2.7: Marlin W4A16 routed-expert MoE kernel (auto-default, Triton fallback) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port the production Marlin INT4 GEMM (the vLLM/sglang serving path for compressed-tensors W4A16) into mstar for the Kimi routed experts, as a clean, model-agnostic quant module. - utils/marlin/: vendored vLLM Marlin CUDA (device code verbatim + classic-ABI host shims), JIT-compiled to torch.ops._mstar_marlin_C.* via the align.py precedent — no vllm/sgl_kernel runtime dep. Repack + fused_marlin_moe launchers. - model/components/quantization/: reusable FusedMoEQuantizeMethod seam, a generic process_weights_after_loading post-load walker, and MarlinMoEMethod. - Kimi wiring: KimiSparseMoeBlock repacks its packed experts to Marlin layout at load and dispatches through it; quant_kernel=auto picks Marlin on sm80+ with the Triton W4A16 path as fallback (quant_kernel=marlin|triton force either). - Tests: kernel + block goldens vs bf16 and Triton (cosine>0.999, relL2<0.02). Validated e2e: real 1T Kimi-K2.7-Code serves at TP8 with Marlin, coherent output. --- .../model/components/quantization/__init__.py | 18 + mstar/model/components/quantization/base.py | 85 + .../components/quantization/marlin_moe.py | 128 + mstar/model/kimi_k2_7/components/moe.py | 150 +- mstar/model/kimi_k2_7/config.py | 42 + mstar/model/kimi_k2_7/kimi_model.py | 6 + mstar/utils/marlin/__init__.py | 10 + mstar/utils/marlin/csrc/core/scalar_type.hpp | 360 +++ .../moe/marlin_moe_wna16/generate_kernels.py | 230 ++ .../moe/marlin_moe_wna16/kernel.h | 47 + .../moe/marlin_moe_wna16/kernel_selector.h | 304 +++ .../moe/marlin_moe_wna16/marlin_moe.cu | 717 ++++++ .../moe/marlin_moe_wna16/marlin_template.h | 2241 +++++++++++++++++ .../sm80_kernel_bfloat16_u4b8_bfloat16.cu | 160 ++ .../sm80_kernel_float16_u4b8_float16.cu | 160 ++ .../quantization/marlin/dequant.h | 609 +++++ .../quantization/marlin/gptq_marlin_repack.cu | 373 +++ .../quantization/marlin/marlin.cuh | 182 ++ .../quantization/marlin/marlin_dtypes.cuh | 149 ++ .../quantization/marlin/marlin_mma.h | 269 ++ mstar/utils/marlin/loader.py | 81 + mstar/utils/marlin/ops.py | 181 ++ mstar/utils/marlin/scalar_type.py | 18 + pyproject.toml | 6 + test/integration/test_kimi_moe_marlin.py | 139 + test/integration/test_marlin_kernels.py | 134 + 26 files changed, 6784 insertions(+), 15 deletions(-) create mode 100644 mstar/model/components/quantization/__init__.py create mode 100644 mstar/model/components/quantization/base.py create mode 100644 mstar/model/components/quantization/marlin_moe.py create mode 100644 mstar/utils/marlin/__init__.py create mode 100644 mstar/utils/marlin/csrc/core/scalar_type.hpp create mode 100644 mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py create mode 100644 mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h create mode 100644 mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_selector.h create mode 100644 mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_moe.cu create mode 100644 mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h create mode 100644 mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_u4b8_bfloat16.cu create mode 100644 mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_float16_u4b8_float16.cu create mode 100644 mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/dequant.h create mode 100644 mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu create mode 100644 mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin.cuh create mode 100644 mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh create mode 100644 mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_mma.h create mode 100644 mstar/utils/marlin/loader.py create mode 100644 mstar/utils/marlin/ops.py create mode 100644 mstar/utils/marlin/scalar_type.py create mode 100644 test/integration/test_kimi_moe_marlin.py create mode 100644 test/integration/test_marlin_kernels.py diff --git a/mstar/model/components/quantization/__init__.py b/mstar/model/components/quantization/__init__.py new file mode 100644 index 000000000..d111578e6 --- /dev/null +++ b/mstar/model/components/quantization/__init__.py @@ -0,0 +1,18 @@ +"""Model-agnostic quantization backends for mstar. + +Currently provides the Marlin W4A16 routed-expert (fused-MoE) path. The seam +(:class:`FusedMoEQuantizeMethod`) and the generic post-load pass +(:func:`process_weights_after_loading`) are model-agnostic; Kimi-K2.7 is the first +consumer. See :mod:`mstar.model.components.quantization.base`. +""" +from mstar.model.components.quantization.base import ( + FusedMoEQuantizeMethod, + process_weights_after_loading, +) +from mstar.model.components.quantization.marlin_moe import MarlinMoEMethod + +__all__ = [ + "FusedMoEQuantizeMethod", + "MarlinMoEMethod", + "process_weights_after_loading", +] diff --git a/mstar/model/components/quantization/base.py b/mstar/model/components/quantization/base.py new file mode 100644 index 000000000..30c4b8e30 --- /dev/null +++ b/mstar/model/components/quantization/base.py @@ -0,0 +1,85 @@ +"""Model-agnostic quantization seams for mstar. + +mstar has no vLLM-style quant-method abstraction. This module introduces the +minimal seam needed to bolt a kernel backend (currently Marlin W4A16 for the +routed experts) onto a model without the model code knowing which kernel runs: + +* :class:`FusedMoEQuantizeMethod` — the interface an MoE block delegates its + quantized routed-expert GEMM to. A block holds one instance, calls + :meth:`~FusedMoEQuantizeMethod.prepare` once post-load to transform its loaded + packed params into the backend's kernel layout, then calls + :meth:`~FusedMoEQuantizeMethod.apply` each forward. Kimi's + :class:`~mstar.model.components.quantization.marlin_moe.MarlinMoEMethod` is the + first implementation; another MoE model can reuse it verbatim. + +* :func:`process_weights_after_loading` — a generic post-load pass. mstar builds + a module on ``meta``, ``to_empty``\\s it, and loads weights, but has no hook to + finalize a kernel layout on the real device afterwards (Marlin needs a one-time + repack + workspace alloc). This walker calls a ``process_weights_after_loading`` + method on every submodule that exposes one; it is a no-op for a plain bf16 model. +""" +from __future__ import annotations + +from typing import Protocol, runtime_checkable + +import torch +from torch import nn + + +@runtime_checkable +class FusedMoEQuantizeMethod(Protocol): + """Kernel backend for a routed-expert (fused-MoE) W4A16 GEMM. + + Implementations own their kernel-format weights after :meth:`prepare` and are + otherwise stateless. The MoE block that holds the method is responsible for + freeing the source packed params once :meth:`prepare` has consumed them. + """ + + def prepare( + self, + w13_packed: torch.Tensor, + w13_scale: torch.Tensor, + w2_packed: torch.Tensor, + w2_scale: torch.Tensor, + device: torch.device, + ) -> None: + """Transform the loaded compressed-tensors packed expert weights into the + backend's runtime layout (e.g. Marlin repack), storing them internally. + + ``w13_packed``/``w2_packed`` are int32 ``(E, N, K // pack_factor)`` and + ``w13_scale``/``w2_scale`` are ``(E, N, K // group_size)`` — the layout + Kimi's Hook B packed params already carry. + """ + ... + + def apply( + self, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + activation: str = "silu", + reduce_results: bool = True, + ) -> torch.Tensor: + """Run the routed-expert GEMM on the prepared weights. + + Mirrors :func:`mstar.utils.fused_moe.fused_experts`: returns + ``(tokens, hidden)`` when ``reduce_results`` else the per-slot + ``(tokens, top_k, hidden)`` tensor the TP path all-reduces before folding. + """ + ... + + +def process_weights_after_loading(root: nn.Module, device: torch.device) -> None: + """Finalize kernel layouts across a freshly-loaded module tree. + + Call once after ``load_weights`` and before ``eval()``/CUDA-graph capture. Any + submodule exposing a ``process_weights_after_loading(device)`` method gets it + invoked (e.g. a Marlin MoE block repacks its packed experts + allocates a + workspace). Modules without the method are skipped, so this is a no-op for a + plain bf16 model. + """ + for module in root.modules(): + hook = getattr(module, "process_weights_after_loading", None) + if callable(hook): + hook(device) diff --git a/mstar/model/components/quantization/marlin_moe.py b/mstar/model/components/quantization/marlin_moe.py new file mode 100644 index 000000000..2400220d8 --- /dev/null +++ b/mstar/model/components/quantization/marlin_moe.py @@ -0,0 +1,128 @@ +"""Marlin W4A16 backend for the routed-expert (fused-MoE) GEMM. + +Implements :class:`~mstar.model.components.quantization.base.FusedMoEQuantizeMethod` +for symmetric INT4 (compressed-tensors ``pack-quantized``, group-wise). It consumes +the packed params an MoE block already loaded (the "Hook B" layout Kimi uses — +``(E, N, K // pack_factor)`` int32 weights + ``(E, N, K // group_size)`` bf16 group +scales), repacks them once into Marlin's tiled layout, and thereafter runs the +vendored Marlin kernels (:mod:`mstar.utils.marlin`). + +Model-agnostic: any MoE block with the same packed-param convention can hold one. +""" +from __future__ import annotations + +import torch + +from mstar.utils.marlin import ops as marlin_ops + + +class MarlinMoEMethod: + """Marlin routed-expert GEMM backend (symmetric INT4, group-wise). + + Stateful: :meth:`prepare` repacks the loaded packed experts into Marlin + layout and stores them (plus the workspace) on the instance; the source + packed params can then be freed by the owning block. :meth:`apply` runs the + two Marlin GEMMs. + """ + + def __init__(self, *, num_bits: int = 4, group_size: int = 32) -> None: + if num_bits != 4: + raise ValueError(f"MarlinMoEMethod supports INT4 only, got num_bits={num_bits}") + self.num_bits = num_bits + self.group_size = group_size + self.pack_factor = 32 // num_bits + # Populated by prepare(): + self.w13_qweight: torch.Tensor | None = None + self.w2_qweight: torch.Tensor | None = None + self.w13_scale: torch.Tensor | None = None + self.w2_scale: torch.Tensor | None = None + self.workspace: torch.Tensor | None = None + + def prepare( + self, + w13_packed: torch.Tensor, + w13_scale: torch.Tensor, + w2_packed: torch.Tensor, + w2_scale: torch.Tensor, + device: torch.device, + ) -> None: + """Repack Hook-B packed experts into Marlin layout. + + mstar packs experts along the input axis as ``(E, N_out, K_in // pack)`` + (compressed-tensors), whereas Marlin's ``gptq_marlin_moe_repack`` wants + GPTQ-style ``(E, K_in // pack, N_out)`` — hence the transpose before each + repack. Scales are permuted the same way. ``w13`` is the fused gate+up + (``N_out = 2 * shard_inter``, ``K_in = hidden``); ``w2`` is the down + projection (``N_out = hidden``, ``K_in = shard_inter``). + """ + pf, gs = self.pack_factor, self.group_size + E, two_inter, hidden_over_pack = w13_packed.shape + hidden = hidden_over_pack * pf + _, w2_hidden, inter_over_pack = w2_packed.shape + inter = inter_over_pack * pf + assert w2_hidden == hidden, f"w2 dim1 {w2_hidden} != hidden {hidden}" + + # gate_up: (E, 2*inter, hidden/pack) -> (E, hidden/pack, 2*inter) -> marlin. + w13_t = w13_packed.transpose(1, 2).contiguous() + self.w13_qweight = marlin_ops.gptq_marlin_moe_repack( + w13_t, size_k=hidden, size_n=two_inter, num_bits=self.num_bits + ) + w13_s = w13_scale.transpose(1, 2).contiguous() # (E, hidden/gs, 2*inter) + self.w13_scale = marlin_ops.marlin_moe_permute_scales( + w13_s, size_k=hidden, size_n=two_inter, group_size=gs + ) + + # down: (E, hidden, inter/pack) -> (E, inter/pack, hidden) -> marlin. + w2_t = w2_packed.transpose(1, 2).contiguous() + self.w2_qweight = marlin_ops.gptq_marlin_moe_repack( + w2_t, size_k=inter, size_n=hidden, num_bits=self.num_bits + ) + w2_s = w2_scale.transpose(1, 2).contiguous() # (E, inter/gs, hidden) + self.w2_scale = marlin_ops.marlin_moe_permute_scales( + w2_s, size_k=inter, size_n=hidden, group_size=gs + ) + + self.workspace = marlin_ops.marlin_make_workspace(torch.device(device)) + + def apply( + self, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + activation: str = "silu", + reduce_results: bool = True, + ) -> torch.Tensor: + assert self.w13_qweight is not None, "MarlinMoEMethod.apply before prepare()" + return marlin_ops.fused_marlin_moe( + x, + self.w13_qweight, + self.w2_qweight, + self.w13_scale, + self.w2_scale, + topk_weights, + topk_ids, + self.workspace, + activation=activation, + reduce_results=reduce_results, + ) + + @staticmethod + def shapes_are_legal(hidden: int, shard_inter: int, group_size: int) -> bool: + """Whether the per-rank expert GEMM shapes satisfy Marlin's tile rules. + + Marlin needs ``n % 64 == 0`` and ``k % 128 == 0`` on each GEMM, plus + ``k % group_size == 0``. gate_up: k=hidden, n=2*shard_inter; down: + k=shard_inter, n=hidden. + """ + if group_size not in (-1, 32, 64, 128): + return False + checks = [ + hidden % 128 == 0, # gate_up K + (2 * shard_inter) % 64 == 0, # gate_up N + shard_inter % 128 == 0, # down K + hidden % 64 == 0, # down N + ] + if group_size != -1: + checks += [hidden % group_size == 0, shard_inter % group_size == 0] + return all(checks) diff --git a/mstar/model/kimi_k2_7/components/moe.py b/mstar/model/kimi_k2_7/components/moe.py index 47a26c375..c7f8e9174 100644 --- a/mstar/model/kimi_k2_7/components/moe.py +++ b/mstar/model/kimi_k2_7/components/moe.py @@ -21,6 +21,8 @@ """ from __future__ import annotations +import logging + import torch import torch.nn.functional as F from torch import nn @@ -35,6 +37,12 @@ ) from mstar.model.kimi_k2_7.config import KimiK2Config +logger = logging.getLogger(__name__) + +# Log the resolved routed-expert backend once per process (the block is +# instantiated per MoE layer, so a per-block log would repeat ~60x). +_BACKEND_LOGGED = False + # --------------------------------------------------------------------------- # Packed-expert weight loaders (int32 weights + bf16 group scales). # @@ -264,6 +272,13 @@ def __init__( self.packed_experts = ( config.quantization_config is not None and config.moe_in_kernel_dequant ) + # Routed-expert W4A16 kernel backend for the packed experts. Marlin layers + # on top of the same packed params; the marlin-vs-triton choice is resolved + # post-load in :meth:`process_weights_after_loading` (a real device is needed + # to probe GPU capability + JIT-build the kernel — ``__init__`` runs on meta). + self.quant_kernel = getattr(config, "quant_kernel", "auto") + self._marlin_method = None + self._use_marlin = False self.gate = KimiMoEGate( hidden_size=config.hidden_size, @@ -289,6 +304,7 @@ def __init__( qc = config.quantization_config self.group_size = qc.group_size self.pack_factor = qc.pack_factor # 8 for INT4 + self.symmetric = qc.symmetric hidden, gs, pf = config.hidden_size, self.group_size, self.pack_factor # The packed/group axes must divide evenly on BOTH the hidden (gate_up K) # and the per-rank intermediate stripe (down K, TP-sharded). @@ -396,22 +412,27 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # Router is replicated: every rank computes the full top-k selection. topk_weights, topk_ids = self.gate(flat) - topk_weights = topk_weights.to(flat.dtype) - if self.packed_experts: - # Packed experts: bypass the shared bf16 ``_dispatch`` and run the - # W4A16 in-kernel dequant GEMM directly (handles tp=1 and tp>1). - routed = self._dispatch_packed_experts(flat, topk_weights, topk_ids) - elif self.tp_size == 1: - routed = _dispatch( - flat, - self.experts.gate_up_proj, - self.experts.down_proj, - self.num_experts, - topk_ids, - topk_weights, - ) + if self._use_marlin: + # Marlin's GEMM takes fp32 combine weights — pass them BEFORE the bf16 + # cast the other paths need. Otherwise identical TP story. + routed = self._dispatch_marlin(flat, topk_weights, topk_ids) else: - routed = self._dispatch_tp(flat, topk_weights, topk_ids) + topk_weights = topk_weights.to(flat.dtype) + if self.packed_experts: + # Packed experts: bypass the shared bf16 ``_dispatch`` and run the + # W4A16 in-kernel dequant GEMM directly (handles tp=1 and tp>1). + routed = self._dispatch_packed_experts(flat, topk_weights, topk_ids) + elif self.tp_size == 1: + routed = _dispatch( + flat, + self.experts.gate_up_proj, + self.experts.down_proj, + self.num_experts, + topk_ids, + topk_weights, + ) + else: + routed = self._dispatch_tp(flat, topk_weights, topk_ids) # Shared expert is a ParallelGatedMLP on the same comm group: at tp>1 it # holds its own intermediate stripe and all-reduces inside its down_proj. shared = self.shared_expert(flat) @@ -456,6 +477,105 @@ def _dispatch_packed_experts( moe_sum_reduce_triton(out, output, routed_scaling_factor=1.0) return output + def process_weights_after_loading(self, device) -> None: + """Resolve the routed-expert kernel backend and, for Marlin, repack the + loaded packed experts into Marlin layout (freeing the source packed params). + + Invoked by the generic post-load walker + (:func:`mstar.model.components.quantization.process_weights_after_loading`) + on a real device — a no-op unless the experts are packed and Marlin is both + selected (``quant_kernel != "triton"``) and eligible (sm80+, symmetric INT4, + Marlin-legal shapes). ``quant_kernel="marlin"`` raises if ineligible so an + explicit request never silently downgrades to Triton. + """ + if not self.packed_experts: + return + from mstar.model.components.quantization import MarlinMoEMethod + from mstar.utils.marlin import is_marlin_available + + dev = torch.device(device) + shard_inter = divide(self.moe_intermediate_size, self.tp_size) + legal_shapes = MarlinMoEMethod.shapes_are_legal( + self.hidden_size, shard_inter, self.group_size + ) + eligible = ( + self.quant_kernel != "triton" + and dev.type == "cuda" + and torch.cuda.get_device_capability(dev) >= (8, 0) + and self.symmetric + and legal_shapes + and is_marlin_available() + ) + if self.quant_kernel == "marlin" and not eligible: + raise RuntimeError( + "quant_kernel='marlin' requested but Marlin is ineligible " + f"(needs CUDA sm80+, symmetric INT4, legal shapes: hidden=" + f"{self.hidden_size}, shard_inter={shard_inter}, " + f"group_size={self.group_size}, legal={legal_shapes}). " + "Use quant_kernel='auto' to fall back to the Triton path." + ) + global _BACKEND_LOGGED + if not eligible: + if not _BACKEND_LOGGED: + logger.info( + "KimiSparseMoeBlock routed-expert backend: Triton W4A16 " + "(quant_kernel=%s, marlin ineligible: legal_shapes=%s).", + self.quant_kernel, legal_shapes, + ) + _BACKEND_LOGGED = True + return + + if not _BACKEND_LOGGED: + logger.info( + "KimiSparseMoeBlock routed-expert backend: Marlin W4A16 " + "(quant_kernel=%s, group_size=%d, tp_size=%d).", + self.quant_kernel, self.group_size, self.tp_size, + ) + _BACKEND_LOGGED = True + + method = MarlinMoEMethod(num_bits=32 // self.pack_factor, group_size=self.group_size) + method.prepare( + self.experts.gate_up_proj_packed.data, + self.experts.gate_up_proj_scale.data, + self.experts.down_proj_packed.data, + self.experts.down_proj_scale.data, + dev, + ) + # Free the source packed params — Marlin holds the repacked copies now. + for name in ( + "gate_up_proj_packed", "gate_up_proj_scale", + "down_proj_packed", "down_proj_scale", + ): + p = getattr(self.experts, name) + p.data = torch.empty(0, dtype=p.dtype, device=dev) + self._marlin_method = method + self._use_marlin = True + + def _dispatch_marlin( + self, + flat: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + ) -> torch.Tensor: + """Marlin W4A16 routed dispatch. TP story is identical to + :meth:`_dispatch_packed_experts`: tp=1 sum-reduces inside the kernel; tp>1 + keeps per-slot partials (``reduce_results=False``), all-reduces the + intermediate-parallel partials, then folds the top-k dim. ``topk_weights`` + is fp32 (Marlin requirement) and already carries ``routed_scaling_factor``. + """ + from mstar.utils.fused_moe import moe_sum_reduce_triton + + reduce = self.tp_size == 1 + out = self._marlin_method.apply( + flat, topk_weights, topk_ids, reduce_results=reduce + ) + if reduce: + return out + self.comm_group.all_reduce(out) + output = torch.empty_like(flat) + moe_sum_reduce_triton(out, output, routed_scaling_factor=1.0) + return output + def _dispatch_tp( self, flat: torch.Tensor, diff --git a/mstar/model/kimi_k2_7/config.py b/mstar/model/kimi_k2_7/config.py index e4041c38f..fc2b51941 100644 --- a/mstar/model/kimi_k2_7/config.py +++ b/mstar/model/kimi_k2_7/config.py @@ -102,6 +102,20 @@ class KimiK2Config: # ``components/moe.py`` / ``weight_loader.py``. moe_in_kernel_dequant: bool = False + # -- Quantization: routed-expert W4A16 kernel backend ---------------------- + # Chooses the kernel for the PACKED routed experts (only meaningful when + # ``moe_in_kernel_dequant`` is set — Marlin layers on top of the Hook B packed + # params). Values: + # "auto" => Marlin on sm80+ (Ampere/Hopper) when the build succeeds and the + # shapes/group_size are Marlin-legal, else the Triton + # ``fused_moe_kernel_w4a16`` fallback. This is the production default. + # "marlin" => force Marlin; raise if ineligible (must not silently downgrade). + # "triton" => force the Triton in-kernel dequant path (the pre-Marlin behavior). + # The final resolution happens post-load in + # ``KimiSparseMoeBlock.process_weights_after_loading`` (a real device is needed to + # probe capability + build the kernel); ``__init__`` runs on ``meta``. + quant_kernel: str = "auto" + # -- Serving: CUDA-graph prefill capture grid (optional overrides) ------ # ``None`` => ``KimiLLMSubmodule`` uses its full-size class-default grid. # ``reduced()`` sets a tiny grid so the synthetic bring-up serve captures a @@ -220,6 +234,34 @@ def reduced_quantized_inkernel( cfg.moe_in_kernel_dequant = True return cfg + @classmethod + def reduced_marlin( + cls, + num_bits: int = 4, + group_size: int = 32, + symmetric: bool = True, + ) -> "KimiK2Config": + """:meth:`reduced_quantized_inkernel` with Marlin-legal shapes + the Marlin + routed-expert backend forced on. + + Marlin's GEMM imposes ``n % 64 == 0`` and ``k % 128 == 0`` on each expert + matmul, which the default reduced dims (``hidden_size=128``, + ``moe_intermediate_size=64``) do NOT satisfy for the down projection + (``k == shard_inter``). This variant bumps ``hidden_size=256`` and + ``moe_intermediate_size=256`` so both expert GEMMs are Marlin-legal at + tp<=2 (tp=1 ``shard_inter=256``, tp=2 ``shard_inter=128`` — both ``% 128``); + tp=4 (``shard_inter=64``) fails ``k % 128``, so pin Marlin goldens to tp<=2. + ``group_size=32`` and ``pack_factor=8`` still divide both axes. + """ + cfg = cls.reduced_quantized_inkernel( + num_bits=num_bits, group_size=group_size, symmetric=symmetric, + ) + cfg.hidden_size = 256 + cfg.moe_intermediate_size = 256 + cfg.intermediate_size = 512 + cfg.quant_kernel = "marlin" + return cfg + @classmethod def k27_code(cls) -> "KimiK2Config": """Full-size ``moonshotai/Kimi-K2.7-Code`` text-only serve config. diff --git a/mstar/model/kimi_k2_7/kimi_model.py b/mstar/model/kimi_k2_7/kimi_model.py index 3e0b9ae77..331272944 100644 --- a/mstar/model/kimi_k2_7/kimi_model.py +++ b/mstar/model/kimi_k2_7/kimi_model.py @@ -406,6 +406,12 @@ def _create_submodule( language_model = language_model.to(autocast_dtype) language_model.to_empty(device=device) load_weights(language_model, source, device=device) + # Post-load pass: let quantized submodules finalize their kernel layout on + # the real device (the routed-expert MoE block repacks its packed experts + # into Marlin layout + allocates a workspace). No-op for a plain bf16 build. + from mstar.model.components.quantization import process_weights_after_loading + + process_weights_after_loading(language_model, torch.device(device)) language_model.eval() logger.info("Successfully loaded Kimi-K2.7 submodule for %s", node_name) diff --git a/mstar/utils/marlin/__init__.py b/mstar/utils/marlin/__init__.py new file mode 100644 index 000000000..fdcd72723 --- /dev/null +++ b/mstar/utils/marlin/__init__.py @@ -0,0 +1,10 @@ +"""Vendored Marlin W4A16 (INT4) CUDA kernels for mstar. + +JIT-compiled from Apache-2.0 vLLM sources under ``csrc/`` on first use, with a +Triton fallback (see :mod:`mstar.utils.marlin.loader`). Exposes the repack + GEMM +launchers (:mod:`mstar.utils.marlin.ops`) used by the compressed-tensors W4A16 +routed-expert path. +""" +from mstar.utils.marlin.loader import is_marlin_available + +__all__ = ["is_marlin_available"] diff --git a/mstar/utils/marlin/csrc/core/scalar_type.hpp b/mstar/utils/marlin/csrc/core/scalar_type.hpp new file mode 100644 index 000000000..b6f39ed79 --- /dev/null +++ b/mstar/utils/marlin/csrc/core/scalar_type.hpp @@ -0,0 +1,360 @@ +#pragma once + +#include +#include +#include +#include +#include + +// For STD_TORCH_CHECK +#include + +namespace vllm { + +// +// ScalarType can represent a wide range of floating point and integer types, +// in particular it can be used to represent sub-byte data types (something +// that torch.dtype currently does not support). +// +// The type definitions on the Python side can be found in: vllm/scalar_type.py +// these type definitions should be kept up to date with any Python API changes +// here. +// +class ScalarType { + public: + enum NanRepr : uint8_t { + NAN_NONE = 0, // nans are not supported + NAN_IEEE_754 = 1, // nans are: exp all 1s, mantissa not all 0s + NAN_EXTD_RANGE_MAX_MIN = 2, // nans are: exp all 1s, mantissa all 1s + + NAN_REPR_ID_MAX + }; + + constexpr ScalarType(uint8_t exponent, uint8_t mantissa, bool signed_, + int32_t bias, bool finite_values_only = false, + NanRepr nan_repr = NAN_IEEE_754) + : exponent(exponent), + mantissa(mantissa), + signed_(signed_), + bias(bias), + finite_values_only(finite_values_only), + nan_repr(nan_repr) {}; + + static constexpr ScalarType int_(uint8_t size_bits, int32_t bias = 0) { + return ScalarType(0, size_bits - 1, true, bias); + } + + static constexpr ScalarType uint(uint8_t size_bits, int32_t bias = 0) { + return ScalarType(0, size_bits, false, bias); + } + + // IEEE 754 compliant floating point type + static constexpr ScalarType float_IEEE754(uint8_t exponent, + uint8_t mantissa) { + STD_TORCH_CHECK(mantissa > 0 && exponent > 0); + return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754); + } + + // IEEE 754 non-compliant floating point type + static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa, + bool finite_values_only, + NanRepr nan_repr) { + STD_TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr"); + STD_TORCH_CHECK(mantissa > 0 && exponent > 0); + STD_TORCH_CHECK( + nan_repr != NAN_IEEE_754, + "use `float_IEEE754` constructor for floating point types that " + "follow IEEE 754 conventions"); + return ScalarType(exponent, mantissa, true, 0, finite_values_only, + nan_repr); + } + + uint8_t const exponent; // size of the exponent field (0 for integer types) + uint8_t const mantissa; // size of the mantissa field (size of the integer + // excluding the sign bit for integer types) + bool const signed_; // flag if the type supports negative numbers (i.e. has a + // sign bit) + int32_t const bias; // stored values equal value + bias, + // used for quantized type + + // Extra Floating point info + bool const finite_values_only; // i.e. no +/-inf if true + NanRepr const nan_repr; // how NaNs are represented + // (not applicable for integer types) + + using Id = int64_t; + + private: + // Field size in id + template + static constexpr size_t member_id_field_width() { + using T = std::decay_t; + return std::is_same_v ? 1 : sizeof(T) * 8; + } + + template + static constexpr auto reduce_members_helper(Fn f, Init val, Member member, + Rest... rest) { + auto new_val = f(val, member); + if constexpr (sizeof...(rest) > 0) { + return reduce_members_helper(f, new_val, rest...); + } else { + return new_val; + }; + } + + template + constexpr auto reduce_members(Fn f, Init init) const { + // Should be in constructor order for `from_id` + return reduce_members_helper(f, init, exponent, mantissa, signed_, bias, + finite_values_only, nan_repr); + }; + + template + static constexpr auto reduce_member_types(Fn f, Init init) { + constexpr auto dummy_type = ScalarType(0, 0, false, 0, false, NAN_NONE); + return dummy_type.reduce_members(f, init); + }; + + static constexpr auto id_size_bits() { + return reduce_member_types( + [](int acc, auto member) -> int { + return acc + member_id_field_width(); + }, + 0); + } + + public: + // unique id for this scalar type that can be computed at compile time for + // c++17 template specialization this is not needed once we migrate to + // c++20 and can pass literal classes as template parameters + constexpr Id id() const { + static_assert(id_size_bits() <= sizeof(Id) * 8, + "ScalarType id is too large to be stored"); + + auto or_and_advance = [](std::pair result, + auto member) -> std::pair { + auto [id, bit_offset] = result; + auto constexpr bits = member_id_field_width(); + return {id | (int64_t(member) & ((uint64_t(1) << bits) - 1)) + << bit_offset, + bit_offset + bits}; + }; + return reduce_members(or_and_advance, std::pair{}).first; + } + + // create a ScalarType from an id, for c++17 template specialization, + // this is not needed once we migrate to c++20 and can pass literal + // classes as template parameters + static constexpr ScalarType from_id(Id id) { + auto extract_and_advance = [id](auto result, auto member) { + using T = decltype(member); + auto [tuple, bit_offset] = result; + auto constexpr bits = member_id_field_width(); + auto extracted_val = static_cast((int64_t(id) >> bit_offset) & + ((uint64_t(1) << bits) - 1)); + auto new_tuple = std::tuple_cat(tuple, std::make_tuple(extracted_val)); + return std::pair{new_tuple, bit_offset + bits}; + }; + + auto [tuple_args, _] = reduce_member_types(extract_and_advance, + std::pair, int>{}); + return std::apply([](auto... args) { return ScalarType(args...); }, + tuple_args); + } + + constexpr int64_t size_bits() const { + return mantissa + exponent + is_signed(); + } + constexpr bool is_signed() const { return signed_; } + constexpr bool is_integer() const { return exponent == 0; } + constexpr bool is_floating_point() const { return exponent > 0; } + constexpr bool is_ieee_754() const { + return is_floating_point() && finite_values_only == false && + nan_repr == NAN_IEEE_754; + } + constexpr bool has_nans() const { + return is_floating_point() && nan_repr != NAN_NONE; + } + constexpr bool has_infs() const { + return is_floating_point() && finite_values_only == false; + } + constexpr bool has_bias() const { return bias != 0; } + + private: + double _floating_point_max() const { + STD_TORCH_CHECK(mantissa <= 52 && exponent <= 11, + "Cannot represent max/min as a double for type ", str()); + + uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1; + if (nan_repr == NAN_EXTD_RANGE_MAX_MIN) { + max_mantissa -= 1; + } + + uint64_t max_exponent = (uint64_t(1) << exponent) - 2; + if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) { + STD_TORCH_CHECK(exponent < 11, + "Cannot represent max/min as a double for type ", str()); + max_exponent += 1; + } + + // adjust the exponent to match that of a double + // for now we assume the exponent bias is the standard 2^(e-1) -1, (where e + // is the exponent bits), there is some precedent for non-standard biases, + // example `float8_e4m3b11fnuz` here: https://github.com/jax-ml/ml_dtypes + // but to avoid premature over complication we are just assuming the + // standard exponent bias until there is a need to support non-standard + // biases + uint64_t exponent_bias = (uint64_t(1) << (exponent - 1)) - 1; + uint64_t exponent_bias_double = (uint64_t(1) << 10) - 1; // double e = 11 + + uint64_t max_exponent_double = + max_exponent - exponent_bias + exponent_bias_double; + + // shift the mantissa into the position for a double and + // the exponent + uint64_t double_raw = + (max_mantissa << (52 - mantissa)) | (max_exponent_double << 52); + + return *reinterpret_cast(&double_raw); + } + + constexpr std::variant _raw_max() const { + if (is_floating_point()) { + return {_floating_point_max()}; + } else { + STD_TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(), + "Cannot represent max as a int64_t"); + return {(int64_t(1) << mantissa) - 1}; + } + } + + constexpr std::variant _raw_min() const { + if (is_floating_point()) { + STD_TORCH_CHECK( + is_signed(), + "We currently assume all floating point types are signed"); + constexpr uint64_t sign_bit_double = (uint64_t(1) << 63); + + double max = _floating_point_max(); + uint64_t max_raw = *reinterpret_cast(&max); + uint64_t min_raw = max_raw | sign_bit_double; + return {*reinterpret_cast(&min_raw)}; + } else { + STD_TORCH_CHECK(!is_signed() || size_bits() <= 64, + "Cannot represent min as a int64_t"); + if (is_signed()) { + // set the top bit to 1 (i.e. INT64_MIN) and the rest to 0 + // then perform an arithmetic shift right to set all the bits above + // (size_bits() - 1) to 1 + return {INT64_MIN >> (64 - size_bits())}; + } else { + return {int64_t(0)}; + } + } + } + + public: + // Max representable value for this scalar type. + // (accounting for bias if there is one) + constexpr std::variant max() const { + return std::visit( + [this](auto x) -> std::variant { return {x - bias}; }, + _raw_max()); + } + + // Min representable value for this scalar type. + // (accounting for bias if there is one) + constexpr std::variant min() const { + return std::visit( + [this](auto x) -> std::variant { return {x - bias}; }, + _raw_min()); + } + + std::string str() const { + /* naming generally follows: https://github.com/jax-ml/ml_dtypes + * for floating point types (leading f) the scheme is: + * `float_em[flags]` + * flags: + * - no-flags: means it follows IEEE 754 conventions + * - f: means finite values only (no infinities) + * - n: means nans are supported (non-standard encoding) + * for integer types the scheme is: + * `[u]int[b]` + * - if bias is not present it means its zero + */ + if (is_floating_point()) { + auto ret = "float" + std::to_string(size_bits()) + "_e" + + std::to_string(exponent) + "m" + std::to_string(mantissa); + if (!is_ieee_754()) { + if (finite_values_only) { + ret += "f"; + } + if (nan_repr != NAN_NONE) { + ret += "n"; + } + } + return ret; + } else { + auto ret = ((is_signed()) ? "int" : "uint") + std::to_string(size_bits()); + if (has_bias()) { + ret += "b" + std::to_string(bias); + } + return ret; + } + } + + constexpr bool operator==(ScalarType const& other) const { + return mantissa == other.mantissa && exponent == other.exponent && + bias == other.bias && signed_ == other.signed_ && + finite_values_only == other.finite_values_only && + nan_repr == other.nan_repr; + } +}; + +using ScalarTypeId = ScalarType::Id; + +// "rust style" names generally following: +// https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L60-L70 +static inline constexpr auto kS4 = ScalarType::int_(4); +static inline constexpr auto kU4 = ScalarType::uint(4); +static inline constexpr auto kU4B8 = ScalarType::uint(4, 8); +static inline constexpr auto kS8 = ScalarType::int_(8); +static inline constexpr auto kU8 = ScalarType::uint(8); +static inline constexpr auto kU8B128 = ScalarType::uint(8, 128); + +static inline constexpr auto kFE2M1f = + ScalarType::float_(2, 1, true, ScalarType::NAN_NONE); +static inline constexpr auto kFE3M2f = + ScalarType::float_(3, 2, true, ScalarType::NAN_NONE); +static inline constexpr auto kFE4M3fn = + ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); +static inline constexpr auto kFE8M0fnu = + ScalarType(8, 0, false, 0, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); +static inline constexpr auto kFE5M2 = ScalarType::float_IEEE754(5, 2); +static inline constexpr auto kFE8M7 = ScalarType::float_IEEE754(8, 7); +static inline constexpr auto kFE5M10 = ScalarType::float_IEEE754(5, 10); + +// Fixed width style names, generally following: +// https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L47-L57 +static inline constexpr auto kInt4 = kS4; +static inline constexpr auto kUint4 = kU4; +static inline constexpr auto kUint4b8 = kU4B8; +static inline constexpr auto kInt8 = kS8; +static inline constexpr auto kUint8 = kU8; +static inline constexpr auto kUint8b128 = kU8B128; + +static inline constexpr auto kFloat4_e2m1f = kFE2M1f; +static inline constexpr auto kFloat6_e3m2f = kFE3M2f; +static inline constexpr auto kFloat8_e4m3fn = kFE4M3fn; +static inline constexpr auto kFloat8_e5m2 = kFE5M2; +static inline constexpr auto kFloat16_e8m7 = kFE8M7; +static inline constexpr auto kFloat16_e5m10 = kFE5M10; + +// colloquial names +static inline constexpr auto kHalf = kFE5M10; +static inline constexpr auto kFloat16 = kHalf; +static inline constexpr auto kBFloat16 = kFE8M7; + +static inline constexpr auto kFloat16Id = kFloat16.id(); +}; // namespace vllm diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py new file mode 100644 index 000000000..3b6fc79f0 --- /dev/null +++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import glob +import itertools +import os +import subprocess +import sys + +import jinja2 + +ARCHS = [] +SUPPORT_FP8 = False +SUPPORT_SM75 = False +SUPPORT_SM80 = False +for arch in sys.argv[1].split(","): + arch = arch[: arch.index(".") + 2].replace(".", "") + arch = int(arch) + # SM89 and the SM12x family (SM120 RTX 5090, SM121 DGX Spark GB10) + # fully support mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32. + # SM90 and SM100 can use this PTX, but it’s simulated + # with FP16 MMA, so it cannot achieve any acceleration. + if arch == 89 or arch // 10 == 12: + SUPPORT_FP8 = True + if arch >= 80: + SUPPORT_SM80 = True + if arch == 75: + SUPPORT_SM75 = True + +FILE_HEAD_COMMENT = """ +// auto generated by generate_kernels.py +// clang-format off +""".lstrip() + +FILE_HEAD = ( + FILE_HEAD_COMMENT + + """ +#include "kernel.h" +#include "marlin_template.h" + +namespace MARLIN_NAMESPACE_NAME { +""" +) + +TEMPLATE = ( + "template __global__ void Marlin<" + "{{a_type_id}}, " + "{{b_type_id}}, " + "{{c_type_id}}, " + "{{s_type_id}}, " + "{{threads}}, " + "{{thread_m_blocks}}, " + "{{thread_n_blocks}}, " + "{{thread_k_blocks}}, " + "{{m_block_size_8}}, " + "{{stages}}, " + "{{group_blocks}}, " + "{{is_zp_float}}>" + "( MARLIN_KERNEL_PARAMS );" +) + +THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)] + +THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4] + +# mstar: trimmed to GPTQ symmetric INT4 (kU4B8) with fp16/bf16 activation only — +# the sole scheme Kimi-K2.7 (compressed-tensors W4A16) uses. All the AWQ / INT8 / +# FP8 / NVFP4 / MXFP4 / W4A8 variants are dropped so the JIT compile stays small. +# Upstream QUANT_CONFIGS kept in git history if another scheme is ever needed. +QUANT_CONFIGS = [ + # GPTQ-INT4 + { + "b_type": "kU4B8", + "thread_configs": THREAD_CONFIGS, + "thread_m_blocks": THREAD_M_BLOCKS, + "group_blocks": [-1, 0, 2, 4, 8], + }, +] + + +def remove_old_kernels(): + for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"): + subprocess.call(["rm", "-f", filename]) + + filename = os.path.dirname(__file__) + "/kernel_selector.h" + subprocess.call(["rm", "-f", filename]) + + +def generate_new_kernels(): + result_dict = {} + sm_75_result_dict = {} + + for quant_config in QUANT_CONFIGS: + c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"]) + a_types = quant_config.get("a_type", ["kFloat16", "kBFloat16"]) + b_type = quant_config["b_type"] + all_group_blocks = quant_config["group_blocks"] + all_m_blocks = quant_config["thread_m_blocks"] + all_thread_configs = quant_config["thread_configs"] + + for a_type, c_type in itertools.product(a_types, c_types): + if not SUPPORT_FP8 and a_type == "kFE4M3fn": + continue + if "16" in a_type and "16" in c_type and a_type != c_type: + continue + s_type = quant_config.get("s_type", c_type) + if (a_type, b_type, c_type) not in result_dict: + result_dict[(a_type, b_type, c_type)] = [] + if a_type in ["kFloat16", "kS8"] and c_type == "kFloat16": + sm_75_result_dict[(a_type, b_type, c_type)] = [] + + for group_blocks, m_blocks, thread_configs in itertools.product( + all_group_blocks, all_m_blocks, all_thread_configs + ): + thread_k, thread_n, threads = thread_configs + + if threads == 256: + # for small batch (m_blocks == 1), + # we only need (128, 128, 256) + # for large batch (m_blocks > 1), + # we only need (64, 256, 256) + if m_blocks <= 1 and (thread_k, thread_n) != (128, 128): + continue + if m_blocks > 1 and (thread_k, thread_n) != (64, 256): + continue + + config = { + "threads": threads, + "s_type": s_type, + "thread_m_blocks": max(m_blocks, 1), + "thread_k_blocks": thread_k // 16, + "thread_n_blocks": thread_n // 16, + "m_block_size_8": "true" if m_blocks == 0.5 else "false", + "stages": 4, + "group_blocks": group_blocks, + "is_zp_float": "false", + } + + if SUPPORT_SM80: + result_dict[(a_type, b_type, c_type)].append(config) + if (a_type, b_type, c_type) in sm_75_result_dict and SUPPORT_SM75: + config_sm75 = config.copy() + config_sm75["stages"] = 2 + sm_75_result_dict[(a_type, b_type, c_type)].append(config_sm75) + + kernel_selector_str = FILE_HEAD_COMMENT + + for result_dict_tmp in [result_dict, sm_75_result_dict]: + for (a_type, b_type, c_type), config_list in result_dict_tmp.items(): + all_template_str_list = [] + if not config_list: + continue + for config in config_list: + s_type = config["s_type"] + template_str = jinja2.Template(TEMPLATE).render( + a_type_id=f"vllm::{a_type}.id()", + b_type_id=f"vllm::{b_type}.id()", + c_type_id=f"vllm::{c_type}.id()", + s_type_id=f"vllm::{s_type}.id()", + **config, + ) + all_template_str_list.append(template_str) + + conditions = [ + f"a_type == vllm::{a_type}", + f"b_type == vllm::{b_type}", + f"c_type == vllm::{c_type}", + f"s_type == vllm::{s_type}", + f"threads == {config['threads']}", + f"thread_m_blocks == {config['thread_m_blocks']}", + f"thread_n_blocks == {config['thread_n_blocks']}", + f"thread_k_blocks == {config['thread_k_blocks']}", + f"m_block_size_8 == {config['m_block_size_8']}", + f"stages == {config['stages']}", + f"group_blocks == {config['group_blocks']}", + f"is_zp_float == {config['is_zp_float']}", + ] + conditions = " && ".join(conditions) + + if kernel_selector_str == FILE_HEAD_COMMENT: + kernel_selector_str += f"if ({conditions})\n kernel = " + else: + kernel_selector_str += f"else if ({conditions})\n kernel = " + + kernel_template2 = ( + "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, " + "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, " + "{{thread_n_blocks}}, {{thread_k_blocks}}, " + "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, " + "{{is_zp_float}}>;" + ) + + kernel_selector_str += ( + jinja2.Template(kernel_template2).render( + a_type_id=f"vllm::{a_type}.id()", + b_type_id=f"vllm::{b_type}.id()", + c_type_id=f"vllm::{c_type}.id()", + s_type_id=f"vllm::{s_type}.id()", + **config, + ) + + "\n" + ) + + file_content = FILE_HEAD + "\n\n" + file_content += "\n\n".join(all_template_str_list) + "\n\n}\n" + if a_type == "kFE4M3fn": + filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" + elif result_dict_tmp is sm_75_result_dict: + filename = f"sm75_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" + else: + filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" + + filename = filename.lower() + + with open(os.path.join(os.path.dirname(__file__), filename), "w") as f: + f.write(file_content) + + if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT: + kernel_selector_str += ( + "else if (a_type == vllm::kFE4M3fn)\n" + " STD_TORCH_CHECK(false, " + '"marlin kernel with fp8 activation is not built.");' + ) + + with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f: + f.write(kernel_selector_str) + + +if __name__ == "__main__": + remove_old_kernels() + generate_new_kernels() diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h new file mode 100644 index 000000000..783736ab5 --- /dev/null +++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h @@ -0,0 +1,47 @@ + +#ifndef MARLIN_NAMESPACE_NAME + #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 +#endif + +#include "libtorch_stable/quantization/marlin/marlin.cuh" +#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" +#include "core/scalar_type.hpp" + +#define MARLIN_KERNEL_PARAMS \ + const int4 *__restrict__ A, const int4 *__restrict__ B, \ + int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ + const int4 *__restrict__ b_bias_ptr, \ + const float *__restrict__ a_scales_ptr, \ + const int4 *__restrict__ scales_ptr, \ + const float *__restrict__ global_scale_ptr, \ + const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \ + const int32_t *__restrict__ sorted_token_ids_ptr, \ + const int32_t *__restrict__ expert_ids_ptr, \ + const int32_t *__restrict__ num_tokens_past_padded_ptr, \ + const float *__restrict__ topk_weights_ptr, int top_k, \ + bool mul_topk_weights, int num_groups, int prob_m, int prob_n, \ + int prob_k, int *locks, bool has_bias, bool use_atomic_add, \ + bool use_fp32_reduce + +namespace MARLIN_NAMESPACE_NAME { +template shared + // fetch pipeline + const int group_blocks, // number of consecutive 16x16 blocks + // with a separate quantization scale + const bool is_zp_float // is zero point of float16 type? + > +__global__ void Marlin(MARLIN_KERNEL_PARAMS); + +} diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_selector.h b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_selector.h new file mode 100644 index 000000000..8551dcd34 --- /dev/null +++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_selector.h @@ -0,0 +1,304 @@ +// auto generated by generate_kernels.py +// clang-format off +if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false) + kernel = Marlin; +else if (a_type == vllm::kFE4M3fn) + STD_TORCH_CHECK(false, "marlin kernel with fp8 activation is not built."); \ No newline at end of file diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_moe.cu b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_moe.cu new file mode 100644 index 000000000..a7452bda1 --- /dev/null +++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_moe.cu @@ -0,0 +1,717 @@ +/* + * Modified by Neural Magic + * Copyright (C) Marlin.2024 Elias Frantar + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * Adapted from https://github.com/IST-DASLab/marlin + */ + +#ifndef MARLIN_NAMESPACE_NAME + #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 +#endif + +#include "kernel.h" + +// mstar: vendored from vLLM. The __global__ kernels + device host helpers +// (marlin_mm etc., lines below) are kept VERBATIM; only the top includes, the +// public host wrapper, and the op registration are ported from vLLM's +// torch::stable ABI to the classic torch ABI mstar's vendored ops use (see +// utils/fused_moe/csrc/moe_align_block_size.cu). STD_TORCH_CHECK, used by the +// verbatim device code, is provided by core/scalar_type.hpp's +// include (reachable via kernel.h). +// +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#include +#include +#include +#include + +#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ + static_assert(std::is_same::value || \ + std::is_same::value, \ + "only float16 and bfloat16 is supported"); + +namespace MARLIN_NAMESPACE_NAME { + +__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){}; + +using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); + +// For a given "a" of size [M,K] performs a permutation of the K columns based +// on the given "perm" indices. +template +__global__ void permute_cols_kernel( + int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr, + int4* __restrict__ out_int4_ptr, + const int32_t* __restrict__ sorted_token_ids_ptr, + const int32_t* __restrict__ expert_ids_ptr, + const int32_t* __restrict__ num_tokens_past_padded_ptr, int size_m, + int size_k, int top_k) { + int num_tokens_past_padded = num_tokens_past_padded_ptr[0]; + int num_moe_blocks = div_ceil(num_tokens_past_padded, moe_block_size); + int32_t block_sorted_ids[moe_block_size]; + int block_num_valid_tokens = 0; + int64_t old_expert_id = 0; + int64_t expert_id = 0; + int row_stride = size_k * sizeof(half) / 16; + + auto read_moe_block_data = [&](int block_id) { + block_num_valid_tokens = moe_block_size; + int4* tmp_block_sorted_ids = reinterpret_cast(block_sorted_ids); + for (int i = 0; i < moe_block_size / 4; i++) { + tmp_block_sorted_ids[i] = + ((int4*)sorted_token_ids_ptr)[block_id * moe_block_size / 4 + i]; + } + for (int i = 0; i < moe_block_size; i++) { + if (block_sorted_ids[i] >= size_m * top_k) { + block_num_valid_tokens = i; + break; + }; + } + }; + + auto permute_row = [&](int row) { + int iters = size_k / default_threads; + int rest = size_k % default_threads; + + int in_offset = (row / top_k) * row_stride; + int out_offset = row * row_stride; + + half const* a_row_half = + reinterpret_cast(a_int4_ptr + in_offset); + half* out_half = reinterpret_cast(out_int4_ptr + out_offset); + + int base_k = 0; + + for (int i = 0; i < iters; i++) { + auto cur_k = base_k + threadIdx.x; + int src_pos = perm_int_ptr[cur_k]; + + out_half[cur_k] = a_row_half[src_pos]; + + base_k += default_threads; + } + + if (rest) { + if (threadIdx.x < rest) { + auto cur_k = base_k + threadIdx.x; + int src_pos = perm_int_ptr[cur_k]; + + out_half[cur_k] = a_row_half[src_pos]; + } + } + }; + + for (int index = blockIdx.x; index < num_moe_blocks; index += gridDim.x) { + old_expert_id = expert_id; + int tmp_expert_id = expert_ids_ptr[index]; + if (tmp_expert_id == -1) continue; + expert_id = tmp_expert_id; + perm_int_ptr += (expert_id - old_expert_id) * size_k; + read_moe_block_data(index); + + for (int i = 0; i < block_num_valid_tokens; i++) + permute_row(block_sorted_ids[i]); + } +} + +typedef struct { + int thread_k; + int thread_n; + int num_threads; +} thread_config_t; + +thread_config_t small_batch_thread_configs[] = { + // Ordered by priority + + // thread_k, thread_n, num_threads + {128, 128, 256}, + {64, 128, 128}, + {128, 64, 128}}; + +thread_config_t large_batch_thread_configs[] = { + // Ordered by priority + + // thread_k, thread_n, num_threads + {64, 256, 256}, + {64, 128, 128}, + {128, 64, 128}}; + +typedef struct { + int blocks_per_sm; + thread_config_t tb_cfg; +} exec_config_t; + +int get_scales_cache_size(thread_config_t const& th_config, int prob_m, + int prob_n, int prob_k, int num_bits, int group_size, + bool has_act_order, bool is_k_full, int stages) { + bool cache_scales_chunk = has_act_order && !is_k_full; + + int tb_n = th_config.thread_n; + int tb_k = th_config.thread_k; + + // Get max scale groups per thread-block + int tb_groups; + if (group_size == -1) { + tb_groups = 1; + } else if (group_size == 0) { + tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size + } else { + tb_groups = div_ceil(tb_k, group_size); + } + + if (cache_scales_chunk) { + int load_groups = + tb_groups * stages * 2; // Chunk size is 2x pipeline over dim K + load_groups = max(load_groups, 32); // We load at least 32 scale groups + return load_groups * tb_n * 2; + } else { + int tb_scales = tb_groups * tb_n * 2; + + return tb_scales * stages; + } +} + +int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8, + int thread_m_blocks, int prob_m, int prob_n, + int prob_k, int num_bits, int group_size, + bool has_act_order, bool is_k_full, int has_zp, + int is_zp_float, bool is_a_8bit, int stages) { + int pack_factor = 32 / num_bits; + + // Get B size + int tb_k = th_config.thread_k; + int tb_n = th_config.thread_n; + int tb_m = thread_m_blocks * 16; + + // shm size for block_sorted_ids/rd_block_sorted_ids/block_topk_weights + // both of them requires tb_m * 4 bytes (tb_m * int32 or tb_m * float32) + int sh_block_meta_size = tb_m * 16; + int sh_a_size = stages * (tb_m * tb_k) * (is_a_8bit ? 1 : 2); + int sh_b_size = stages * (tb_k * tb_n / pack_factor) * 4; + int sh_red_size = tb_m * (tb_n + 8) * 2; + int sh_bias_size = tb_n * 2; + int tmp_size = + (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size; + tmp_size = max(max(sh_b_size, sh_red_size), tmp_size); + + int sh_s_size = + get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, + group_size, has_act_order, is_k_full, stages); + int sh_g_idx_size = has_act_order && !is_k_full ? stages * tb_k / 4 : 0; + int sh_zp_size = 0; + if (has_zp) { + if (is_zp_float) + sh_zp_size = sh_s_size; + else if (num_bits == 4) + sh_zp_size = sh_s_size / 4; + else if (num_bits == 8) + sh_zp_size = sh_s_size / 2; + } + + int total_size = tmp_size + sh_a_size + sh_s_size + sh_zp_size + + sh_g_idx_size + sh_block_meta_size; + + return total_size; +} + +bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8, + int thread_m_blocks, int prob_m, int prob_n, int prob_k, + int num_bits, int group_size, bool has_act_order, + bool is_k_full, int has_zp, int is_zp_float, + bool is_a_8bit, int stages, int max_shared_mem) { + // Sanity + if (th_config.thread_k == -1 || th_config.thread_n == -1 || + th_config.num_threads == -1) { + return false; + } + + // Verify K/N are divisible by thread K/N + if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) { + return false; + } + + // Verify min for thread K/N + if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) { + return false; + } + + // num_threads must be at least 128 (= 4 warps) + if (th_config.num_threads < 128) { + return false; + } + + // Check that pipeline fits into cache + int cache_size = + get_kernel_cache_size(th_config, m_block_size_8, thread_m_blocks, prob_m, + prob_n, prob_k, num_bits, group_size, has_act_order, + is_k_full, has_zp, is_zp_float, is_a_8bit, stages); + return cache_size <= max_shared_mem; +} + +MarlinFuncPtr get_marlin_kernel( + const vllm::ScalarType a_type, const vllm::ScalarType b_type, + const vllm::ScalarType c_type, const vllm::ScalarType s_type, + int thread_m_blocks, int thread_n_blocks, int thread_k_blocks, + bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks, + int threads, bool is_zp_float, int stages) { + int num_bits = b_type.size_bits(); + auto kernel = MarlinDefault; + +#include "kernel_selector.h" + + return kernel; +} + +exec_config_t determine_exec_config( + const vllm::ScalarType& a_type, const vllm::ScalarType& b_type, + const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m, + int prob_n, int prob_k, int num_experts, int top_k, int thread_m_blocks, + bool m_block_size_8, int num_bits, int group_size, bool has_act_order, + bool is_k_full, bool has_zp, bool is_zp_float, bool is_a_8bit, int stages, + int max_shared_mem, int sms) { + exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}}; + thread_config_t* thread_configs = thread_m_blocks > 1 + ? large_batch_thread_configs + : small_batch_thread_configs; + int thread_configs_size = + thread_m_blocks > 1 + ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t) + : sizeof(small_batch_thread_configs) / sizeof(thread_config_t); + + int count = 0; + constexpr int device_max_reg_size = 255 * 1024; + for (int i = 0; i < thread_configs_size; i++) { + thread_config_t th_config = thread_configs[i]; + + if (!is_valid_config(th_config, m_block_size_8, thread_m_blocks, prob_m, + prob_n, prob_k, num_bits, group_size, has_act_order, + is_k_full, has_zp, is_zp_float, is_a_8bit, stages, + max_shared_mem - 512)) { + continue; + } + + int cache_size = get_kernel_cache_size( + th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k, + num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float, + is_a_8bit, stages); + + int group_blocks = 0; + if (!has_act_order) { + group_blocks = group_size == -1 ? -1 : (group_size / 16); + } + + auto kernel = + get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks, + th_config.thread_n / 16, th_config.thread_k / 16, + m_block_size_8, has_act_order, has_zp, group_blocks, + th_config.num_threads, is_zp_float, stages); + + if (kernel == MarlinDefault) continue; + + cudaFuncAttributes attr; + cudaFuncGetAttributes(&attr, kernel); + int reg_size = max(attr.numRegs, 1) * th_config.num_threads * 4; + int allow_count = min(device_max_reg_size / reg_size, + max_shared_mem / (cache_size + 1536)); + if (thread_m_blocks == 1) + allow_count = max(min(allow_count, 4), 1); + else + allow_count = max(min(allow_count, 2), 1); + + if (prob_n / th_config.thread_n * prob_m * top_k * 4 < sms * allow_count) { + allow_count = + max(prob_n / th_config.thread_n * prob_m * top_k * 4 / sms, 1); + } + + if (allow_count > count) { + count = allow_count; + exec_cfg = {count, th_config}; + }; + } + + return exec_cfg; +} + +void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias, + void* a_s, void* b_s, void* g_s, void* zp, void* g_idx, + void* perm, void* a_tmp, void* sorted_token_ids, + void* expert_ids, void* num_tokens_past_padded, + void* topk_weights, int moe_block_size, int num_experts, + int top_k, bool mul_topk_weights, int prob_m, int prob_n, + int prob_k, void* workspace, vllm::ScalarType const& a_type, + vllm::ScalarType const& b_type, vllm::ScalarType const& c_type, + vllm::ScalarType const& s_type, bool has_bias, + bool has_act_order, bool is_k_full, bool has_zp, int num_groups, + int group_size, int dev, cudaStream_t stream, int thread_k, + int thread_n, int sms, int blocks_per_sm, bool use_atomic_add, + bool use_fp32_reduce, bool is_zp_float) { + int thread_m_blocks = div_ceil(moe_block_size, 16); + bool m_block_size_8 = moe_block_size == 8; + bool is_a_8bit = a_type.size_bits() == 8; + + STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", + prob_m, ", ", prob_n, ", ", prob_k, "]"); + + int group_blocks = 0; + if (has_act_order) { + if (is_k_full) { + STD_TORCH_CHECK(group_size != -1); + group_blocks = group_size / 16; + STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } else { + STD_TORCH_CHECK(group_size == 0); + group_blocks = 0; + } + } else { + if (group_size == -1) { + group_blocks = -1; + } else { + group_blocks = group_size / 16; + STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k, + " is not divisible by group_blocks = ", group_blocks); + } + } + + int num_bits = b_type.size_bits(); + const int4* A_ptr = (const int4*)A; + const int4* B_ptr = (const int4*)B; + int4* C_ptr = (int4*)C; + int4* C_tmp_ptr = (int4*)C_tmp; + const int4* bias_ptr = (const int4*)b_bias; + const float* a_s_ptr = (const float*)a_s; + const int4* b_s_ptr = (const int4*)b_s; + const float* g_s_ptr = (const float*)g_s; + const int4* zp_ptr = (const int4*)zp; + const int* g_idx_ptr = (const int*)g_idx; + const int* perm_ptr = (const int*)perm; + int4* a_tmp_ptr = (int4*)a_tmp; + const int32_t* sorted_token_ids_ptr = (const int32_t*)sorted_token_ids; + const int32_t* expert_ids_ptr = (const int32_t*)expert_ids; + const int32_t* num_tokens_past_padded_ptr = + (const int32_t*)num_tokens_past_padded; + const float* topk_weights_ptr = (const float*)topk_weights; + int* locks = (int*)workspace; + + if (has_act_order) { + // Permute A columns + auto kernel = permute_cols_kernel<8>; + if (moe_block_size == 8) { + } else if (moe_block_size == 16) + kernel = permute_cols_kernel<16>; + else if (moe_block_size == 32) + kernel = permute_cols_kernel<32>; + else if (moe_block_size == 48) + kernel = permute_cols_kernel<48>; + else if (moe_block_size == 64) + kernel = permute_cols_kernel<64>; + else + STD_TORCH_CHECK(false, "unsupported moe_block_size ", moe_block_size); + + // avoid ">>>" being formatted to "> > >" + // clang-format off + kernel<<>>( + A_ptr, perm_ptr, a_tmp_ptr, sorted_token_ids_ptr, expert_ids_ptr, + num_tokens_past_padded_ptr, prob_m, prob_k, top_k); + // clang-format on + A_ptr = a_tmp_ptr; + prob_m = prob_m * top_k; + top_k = 1; + + // If we have a full K, then we can run the non-act-order version of Marlin + // (since the weight rows are reordered by increasing group ids, and by + // having a full K, we have full original groups) + if (is_k_full) has_act_order = false; + } + + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + STD_TORCH_CHECK(max_shared_mem > 0); + + int major_capability, minor_capability; + cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, + dev); + cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, + dev); + STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 75, + "marlin kernel only support Turing or newer GPUs."); + int stages = 4; + if (major_capability == 7 && minor_capability == 5) { + stages = 2; + STD_TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8, + "Turing only support FP16 or INT8 activation."); + } + if (a_type == vllm::kFE4M3fn) { + STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 89, + "FP8 only support Ada Lovelace or newer GPUs."); + STD_TORCH_CHECK( + major_capability * 10 + minor_capability == 89 || + major_capability == 12, + "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than " + "Marlin W4A16 on other devices)."); + } + + // Set thread config + exec_config_t exec_cfg; + thread_config_t thread_tfg; + if (thread_k != -1 && thread_n != -1) { + thread_tfg = thread_config_t{thread_k, thread_n, thread_k * thread_n / 64}; + if (blocks_per_sm == -1) blocks_per_sm = 1; + exec_cfg = exec_config_t{blocks_per_sm, thread_tfg}; + STD_TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, + " is not divisible by thread_n = ", thread_n); + STD_TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, + " is not divisible by thread_k = ", thread_k); + } else { + // Auto config + exec_cfg = determine_exec_config( + a_type, b_type, c_type, s_type, prob_m, prob_n, prob_k, num_experts, + top_k, thread_m_blocks, m_block_size_8, num_bits, group_size, + has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages, + max_shared_mem, sms); + thread_tfg = exec_cfg.tb_cfg; + } + + int num_threads = thread_tfg.num_threads; + thread_k = thread_tfg.thread_k; + thread_n = thread_tfg.thread_n; + int blocks = sms * exec_cfg.blocks_per_sm; + if (exec_cfg.blocks_per_sm > 1) + max_shared_mem = max_shared_mem / exec_cfg.blocks_per_sm - 1024; + + int thread_k_blocks = thread_k / 16; + int thread_n_blocks = thread_n / 16; + + STD_TORCH_CHECK( + is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks, prob_m, + prob_n, prob_k, num_bits, group_size, has_act_order, + is_k_full, has_zp, is_zp_float, is_a_8bit, stages, + max_shared_mem), + "Invalid thread config: thread_m_blocks = ", thread_m_blocks, + ", thread_k = ", thread_tfg.thread_k, + ", thread_n = ", thread_tfg.thread_n, + ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ", + prob_k, ", ", prob_n, "] and num_bits = ", num_bits, + ", group_size = ", group_size, ", has_act_order = ", has_act_order, + ", is_k_full = ", is_k_full, ", has_zp = ", has_zp, + ", is_zp_float = ", is_zp_float, ", max_shared_mem = ", max_shared_mem); + + int sh_cache_size = + get_kernel_cache_size(thread_tfg, m_block_size_8, thread_m_blocks, prob_m, + prob_n, prob_k, num_bits, group_size, has_act_order, + is_k_full, has_zp, is_zp_float, is_a_8bit, stages); + + auto kernel = get_marlin_kernel( + a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks, + thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks, + num_threads, is_zp_float, stages); + + if (kernel == MarlinDefault) { + STD_TORCH_CHECK( + false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ", + prob_k, "]", ", has_act_order = ", has_act_order, + ", num_groups = ", num_groups, ", group_size = ", group_size, + ", thread_m_blocks = ", thread_m_blocks, + ", thread_n_blocks = ", thread_n_blocks, + ", thread_k_blocks = ", thread_k_blocks, ", num_bits = ", num_bits); + } + + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + max_shared_mem); + // avoid ">>>" being formatted to "> > >" + // clang-format off + kernel<<>>( + A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, g_idx_ptr, + sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr, + topk_weights_ptr, top_k, mul_topk_weights, num_groups, prob_m, + prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce); + // clang-format on +} + +} // namespace MARLIN_NAMESPACE_NAME + +// --------------------------------------------------------------------------- +// mstar classic-ABI host wrapper, trimmed to W4A16: fp16/bf16 activation, +// symmetric GPTQ INT4 (uint4b8), no zero-point / act-order / bias / global-scale +// / 8-bit-activation. Ports vLLM ops.cu::moe_wna16_marlin_gemm; every dropped +// optional is handed to the verbatim device ``marlin_mm`` as a 0-element tensor, +// exactly as upstream's "absent optional" branches did. Namespace of the device +// symbols is ``marlin_moe_wna16`` (set by kernel.h's MARLIN_NAMESPACE_NAME), so +// its Marlin<> instantiations never ODR-clash with the linear repack's ``marlin``. +// --------------------------------------------------------------------------- +torch::Tensor moe_wna16_marlin_gemm( + torch::Tensor a, std::optional c_or_none, + torch::Tensor b_q_weight, torch::Tensor b_scales, torch::Tensor workspace, + torch::Tensor sorted_token_ids, torch::Tensor expert_ids, + torch::Tensor num_tokens_past_padded, torch::Tensor topk_weights, + int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, + int64_t b_type_id, int64_t size_m, int64_t size_n, int64_t size_k, + bool is_k_full, bool use_atomic_add, bool use_fp32_reduce) { + vllm::ScalarTypeId a_type_id, c_type_id; + TORCH_CHECK( + a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16, + "moe_wna16_marlin_gemm (mstar W4A16): activation must be fp16 or bf16"); + if (a.scalar_type() == torch::kHalf) { + a_type_id = vllm::kFloat16.id(); + c_type_id = vllm::kFloat16.id(); + } else { + a_type_id = vllm::kBFloat16.id(); + c_type_id = vllm::kBFloat16.id(); + } + auto c_dtype = a.scalar_type(); + vllm::ScalarTypeId s_type_id = c_type_id; + + vllm::ScalarType a_type = vllm::ScalarType::from_id(a_type_id); + vllm::ScalarType b_type = vllm::ScalarType::from_id(b_type_id); + vllm::ScalarType c_type = vllm::ScalarType::from_id(c_type_id); + vllm::ScalarType s_type = vllm::ScalarType::from_id(s_type_id); + TORCH_CHECK(b_type == vllm::kU4B8, + "mstar Marlin MoE supports only symmetric INT4 (uint4b8); got ", + b_type.str()); + + int pack_factor = 32 / b_type.size_bits(); + int num_experts = b_q_weight.size(0); + + if (moe_block_size != 8) { + TORCH_CHECK(moe_block_size % 16 == 0, + "unsupported moe_block_size=", moe_block_size); + TORCH_CHECK(moe_block_size >= 16 && moe_block_size <= 64, + "unsupported moe_block_size=", moe_block_size); + } + + // Verify A + TORCH_CHECK(a.size(0) == size_m, "a.size(0) = ", a.size(0), + ", size_m = ", size_m); + TORCH_CHECK(a.size(1) == size_k, "a.size(1) = ", a.size(1), + ", size_k = ", size_k); + + // Verify B (Marlin-tiled: b_q_weight is (E, size_k/tile, size_n*pack/tile)) + TORCH_CHECK(size_k % marlin_moe_wna16::tile_size == 0, "size_k = ", size_k, + " not divisible by tile_size = ", marlin_moe_wna16::tile_size); + TORCH_CHECK((size_k / marlin_moe_wna16::tile_size) == b_q_weight.size(1), + "b_q_weight.size(1) = ", b_q_weight.size(1), ", size_k = ", size_k); + TORCH_CHECK(b_q_weight.size(2) % marlin_moe_wna16::tile_size == 0, + "b_q_weight.size(2) = ", b_q_weight.size(2)); + int actual_size_n = + (b_q_weight.size(2) / marlin_moe_wna16::tile_size) * pack_factor; + TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n, + ", actual_size_n = ", actual_size_n); + + TORCH_CHECK(a.is_cuda() && a.is_contiguous(), "A must be contiguous CUDA"); + TORCH_CHECK(b_q_weight.is_cuda() && b_q_weight.is_contiguous(), + "b_q_weight must be contiguous CUDA"); + TORCH_CHECK(b_scales.is_cuda() && b_scales.is_contiguous(), + "b_scales must be contiguous CUDA"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); + int dev = a.get_device(); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + int sms = -1; + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + + auto opts_c = a.options().dtype(c_dtype); + auto opts_f = a.options().dtype(torch::kFloat); + + torch::Tensor c; + if (c_or_none.has_value()) { + c = c_or_none.value(); + TORCH_CHECK(c.is_cuda() && c.is_contiguous(), "c must be contiguous CUDA"); + TORCH_CHECK(c.size(0) == size_m * top_k && c.size(1) == size_n, + "bad c shape"); + } else { + c = torch::empty({size_m * top_k, size_n}, opts_c); + } + + torch::Tensor c_tmp; + if (use_fp32_reduce && !use_atomic_add) { + long max_c_tmp_size = std::min( + (long)size_n * sorted_token_ids.size(0), + (long)sms * 4 * moe_block_size * marlin_moe_wna16::max_thread_n); + if (moe_block_size == 8) max_c_tmp_size *= 2; + c_tmp = torch::empty({max_c_tmp_size}, opts_f); + } else { + c_tmp = torch::empty({0}, opts_f); + } + + // Grouping (no act-order): num_groups = b_scales.size(1). + TORCH_CHECK(b_scales.dim() == 3, "b_scales must be rank 3"); + TORCH_CHECK(b_scales.size(2) == size_n, "b_scales dim2 != size_n"); + int num_groups = b_scales.size(1); + int group_size; + if (num_groups > 1) { + TORCH_CHECK(size_k % num_groups == 0, "size_k not divisible by num_groups"); + group_size = size_k / num_groups; + } else { + group_size = -1; + } + + // Dropped optionals -> 0-element tensors (upstream's absent-optional branches). + torch::Tensor a_scales = torch::empty({0}, opts_f); + torch::Tensor global_scale = torch::empty({0}, opts_f); + torch::Tensor b_bias = torch::empty({0}, opts_c); + torch::Tensor b_zeros = torch::empty({0}, opts_c); + torch::Tensor g_idx = torch::empty({0}, opts_c); + torch::Tensor perm = torch::empty({0}, opts_c); + torch::Tensor a_tmp = torch::empty({0}, opts_c); + + TORCH_CHECK(size_n % marlin_moe_wna16::min_thread_n == 0, "size_n = ", size_n, + " not divisible by min_thread_n"); + int max_n_tiles = size_n / marlin_moe_wna16::min_thread_n; + int min_workspace_size = std::min( + max_n_tiles * (int)(sorted_token_ids.size(0) / moe_block_size), sms * 4); + TORCH_CHECK(workspace.numel() >= min_workspace_size, + "workspace.numel = ", workspace.numel(), + " < min_workspace_size = ", min_workspace_size); + + marlin_moe_wna16::marlin_mm( + a.const_data_ptr(), b_q_weight.const_data_ptr(), c.mutable_data_ptr(), + c_tmp.mutable_data_ptr(), b_bias.mutable_data_ptr(), + a_scales.mutable_data_ptr(), b_scales.mutable_data_ptr(), + global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(), + g_idx.mutable_data_ptr(), perm.mutable_data_ptr(), a_tmp.mutable_data_ptr(), + sorted_token_ids.mutable_data_ptr(), expert_ids.mutable_data_ptr(), + num_tokens_past_padded.mutable_data_ptr(), topk_weights.mutable_data_ptr(), + moe_block_size, num_experts, top_k, mul_topk_weights, size_m, size_n, + size_k, workspace.mutable_data_ptr(), a_type, b_type, c_type, s_type, + /*has_bias=*/false, /*has_act_order=*/false, is_k_full, /*has_zp=*/false, + num_groups, group_size, dev, stream, /*thread_k=*/-1, /*thread_n=*/-1, sms, + /*blocks_per_sm=*/-1, use_atomic_add, use_fp32_reduce, + /*is_zp_float=*/false); + + return c; +} + +// Single TORCH_LIBRARY block for the whole ``_mstar_marlin_C`` namespace (the +// repack impl lives in gptq_marlin_repack.cu as a TORCH_LIBRARY_IMPL). +TORCH_LIBRARY(_mstar_marlin_C, m) { + m.def( + "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, int size_k, " + "int size_n, int num_bits) -> Tensor"); + m.def( + "moe_wna16_marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, " + "Tensor b_scales, Tensor workspace, Tensor sorted_token_ids, " + "Tensor expert_ids, Tensor num_tokens_past_padded, Tensor topk_weights, " + "int moe_block_size, int top_k, bool mul_topk_weights, int b_type_id, " + "int size_m, int size_n, int size_k, bool is_k_full, bool use_atomic_add, " + "bool use_fp32_reduce) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(_mstar_marlin_C, CUDA, m) { + m.impl("moe_wna16_marlin_gemm", &moe_wna16_marlin_gemm); +} diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h new file mode 100644 index 000000000..04f90101b --- /dev/null +++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h @@ -0,0 +1,2241 @@ +/* + * Modified by Neural Magic + * Copyright (C) Marlin.2024 Elias Frantar + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * Adapted from https://github.com/IST-DASLab/marlin + */ + +#ifndef MARLIN_NAMESPACE_NAME + #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 +#endif + +#include "libtorch_stable/quantization/marlin/marlin.cuh" +#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" +#include "libtorch_stable/quantization/marlin/dequant.h" +#include "libtorch_stable/quantization/marlin/marlin_mma.h" +#include "core/scalar_type.hpp" + +#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ + static_assert(std::is_same::value || \ + std::is_same::value, \ + "only float16 and bfloat16 is supported"); + +namespace MARLIN_NAMESPACE_NAME { + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750 + +template shared + // fetch pipeline + const bool has_act_order, // whether act_order is enabled + const int group_blocks, // number of consecutive 16x16 blocks + // with a separate quantization scale + const bool is_zp_float // is zero point of float16 type? + > +__global__ void Marlin( + const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape + // (k/groupsize)xn + const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape + // (k/groupsize)x(n/pack_factor) + const int* __restrict__ g_idx, // int32 group indices of shape k + const int32_t* __restrict__ sorted_token_ids_ptr, // moe sorted_ids + const int32_t* __restrict__ expert_ids_ptr, // moe expert ids + const int32_t* __restrict__ num_tokens_past_padded_ptr, // moe num tokens + const float* __restrict__ topk_weights_ptr, // moe top weights + int top_k, // num of experts per token + bool mul_topk_weights, // mul topk weights or not + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks, // extra global storage for barrier synchronization + bool use_atomic_add, // whether to use atomic add to reduce + bool use_fp32_reduce // whether to use fp32 global reduce +) {} + +} // namespace MARLIN_NAMESPACE_NAME + +#else + +// Instruction for loading a full 16x16 matrix fragment of operand A from shared +// memory, directly in tensor core layout. +template +__device__ inline void ldsm(typename MarlinScalarType::FragA& frag_a, + const void* smem_ptr) { + uint32_t* a = reinterpret_cast(&frag_a); + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + if constexpr (count == 4) { + asm volatile( + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(smem)); + } else if constexpr (count == 2) { + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" + : "=r"(a[0]), "=r"(a[1]) + : "r"(smem)); + } else if constexpr (count == 1) { + asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" + : "=r"(a[0]) + : "r"(smem)); + } else { + static_assert(count == 1 || count == 2 || count == 4, "invalid count"); + } +} + +// Multiply dequantized values by the corresponding quantization scale; used +// only for grouped quantization. +template +__device__ inline void scale(typename MarlinScalarType::FragB& frag_b, + typename MarlinScalarType::FragS& frag_s, + int i) { + using scalar_t = typename MarlinScalarType::scalar_t; + using scalar_t2 = typename MarlinScalarType::scalar_t2; + scalar_t2 s = MarlinScalarType::num2num2( + reinterpret_cast(&frag_s)[i]); + frag_b[0] = __hmul2(frag_b[0], s); + frag_b[1] = __hmul2(frag_b[1], s); +} + +template +__device__ inline void scale_and_sub( + typename MarlinScalarType::FragB& frag_b, + typename MarlinScalarType::scalar_t s, + typename MarlinScalarType::scalar_t zp) { + using scalar_t = typename MarlinScalarType::scalar_t; + using scalar_t2 = typename MarlinScalarType::scalar_t2; + scalar_t2 s2 = MarlinScalarType::num2num2(s); + scalar_t2 zp2 = MarlinScalarType::num2num2(zp); + frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2)); + frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2)); +} + +template +__device__ inline void sub_zp( + typename MarlinScalarType::FragB& frag_b, + typename MarlinScalarType::scalar_t2& frag_zp, int i) { + using scalar_t = typename MarlinScalarType::scalar_t; + using scalar_t2 = typename MarlinScalarType::scalar_t2; + scalar_t2 zp = MarlinScalarType::num2num2( + reinterpret_cast(&frag_zp)[i]); + frag_b[0] = __hsub2(frag_b[0], zp); + frag_b[1] = __hsub2(frag_b[1], zp); +} + +// Same as above, but for act_order (each K is multiplied individually) +template +__device__ inline void scale4( + typename MarlinScalarType::FragB& frag_b, + typename MarlinScalarType::FragS& frag_s_1, + typename MarlinScalarType::FragS& frag_s_2, + typename MarlinScalarType::FragS& frag_s_3, + typename MarlinScalarType::FragS& frag_s_4, int i) { + using scalar_t = typename MarlinScalarType::scalar_t; + using scalar_t2 = typename MarlinScalarType::scalar_t2; + + scalar_t2 s_val_1_2; + s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; + s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; + + scalar_t2 s_val_3_4; + s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; + s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; + + frag_b[0] = __hmul2(frag_b[0], s_val_1_2); + frag_b[1] = __hmul2(frag_b[1], s_val_3_4); +} + +// Given 2 floats multiply by 2 scales (halves) +template +__device__ inline void scale_float( + float* c, typename MarlinScalarType::FragS& s) { + using scalar_t = typename MarlinScalarType::scalar_t; + scalar_t* s_ptr = reinterpret_cast(&s); + c[0] = __fmul_rn(c[0], MarlinScalarType::num2float(s_ptr[0])); + c[1] = __fmul_rn(c[1], MarlinScalarType::num2float(s_ptr[1])); +} + +// Wait until barrier reaches `count`, then lock for current threadblock. +__device__ inline void barrier_acquire(int* lock, int count) { + if (threadIdx.x == 0) { + int state = -1; + do + // Guarantee that subsequent writes by this threadblock will be visible + // globally. + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" + : "=r"(state) + : "l"(lock)); + while (state != count); + } + __syncthreads(); +} + +// Release barrier and increment visitation count. +__device__ inline void barrier_release(int* lock, bool reset = false) { + __syncthreads(); + if (threadIdx.x == 0) { + if (reset) { + lock[0] = 0; + return; + } + int val = 1; + // Make sure that all writes since acquiring this barrier are visible + // globally, while releasing the barrier. + asm volatile("fence.acq_rel.gpu;\n"); + asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" + : + : "l"(lock), "r"(val)); + } +} + +// Wait until value of lock to be negative, and then add 1 +__device__ inline void wait_negative_and_add(int* lock) { + if (threadIdx.x == 0) { + int state = 0; + do + // Guarantee that subsequent writes by this threadblock will be visible + // globally. + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" + : "=r"(state) + : "l"(lock)); + while (state >= 0); + atomicAdd(lock, 1); + } + __syncthreads(); +} + +template shared + // fetch pipeline + const int group_blocks, // number of consecutive 16x16 blocks + // with a separate quantization scale + const bool is_zp_float // is zero point of float16 type? + > +__global__ void Marlin( + const int4* __restrict__ A, // fp16 input matrix of shape mxk + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn + int4* __restrict__ C, // fp16 output buffer of shape mxn + int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) + const int4* __restrict__ b_bias_ptr, + // float scales of input matrix, only used when is_a_8bit == true. + // shape (m,) + const float* __restrict__ a_scales_ptr, + // fp16 quantization scales. shape (k/groupsize, n) + const int4* __restrict__ scales_ptr, + // fp16 global scale (for nvfp4// only) + const float* __restrict__ global_scale_ptr, + // 4bit packed zero-points of shape + // (k/groupsize, n/pack_factor) + const int4* __restrict__ zp_ptr, + // int32 group indices of shape k + const int* __restrict__ g_idx, + const int32_t* __restrict__ sorted_token_ids_ptr, // moe sorted_ids + const int32_t* __restrict__ expert_ids_ptr, // moe expert ids + const int32_t* __restrict__ num_tokens_past_padded_ptr, // moe num tokens + const float* __restrict__ topk_weights_ptr, // moe top weights + int top_k, // num of experts per token + bool mul_topk_weights, // mul topk weights or not + int num_groups, // number of scale groups per output channel + int prob_m, // batch dimension m + int prob_n, // output dimension n + int prob_k, // reduction dimension k + int* locks, // extra global storage for barrier synchronization + bool has_bias, + bool use_atomic_add, // whether to use atomic add to reduce + bool use_fp32_reduce // whether to use fp32 global reduce +) { + // Each threadblock processes one "stripe" of the B matrix with (roughly) the + // same size, which might involve multiple column "slices" (of width 16 * + // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM + // example: + // 0 1 3 + // 0 2 3 + // 1 2 4 + // While this kind of partitioning makes things somewhat more complicated, it + // ensures good utilization of all SMs for many kinds of shape and GPU + // configurations, while requiring as few slow global cross-threadblock + // reductions as possible. + + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 890 + // FP8 computation is only supported for Ada Lovelace or newer architectures. + if constexpr (a_type_id == vllm::kFE4M3fn.id()) return; + #endif + + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 + // Turing TensorCore only supports fp16 and int8 + if constexpr (a_type_id != vllm::kFloat16.id() && a_type_id != vllm::kS8.id()) + return; + #endif + + int num_tokens_past_padded = num_tokens_past_padded_ptr[0]; + constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); + + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 + static constexpr auto num_bits = + vllm::ScalarType::from_id(b_type_id).size_bits(); + // Disable use_fp16_accum for NVFP4 and cases when group_size == -1 && + // num_bits == 4 + constexpr bool use_fp16_accum = + a_type_id == vllm::kFloat16.id() && + (!(b_type_id == vllm::kFE2M1f.id() && s_type_id == vllm::kFE4M3fn.id()) && + !(group_blocks == -1 && num_bits == 4)); + #else + constexpr bool use_fp16_accum = false; + #endif + using Adtype = MarlinScalarType; + using Cdtype = MarlinScalarType; + + using scalar_t = typename MarlinScalarType::scalar_t; + using scalar_t2 = typename MarlinScalarType::scalar_t2; + using scalar_32bit_t = typename MarlinScalarType::scalar_32bit_t; + + using c_scalar_t = typename MarlinScalarType::scalar_t; + using c_scalar_t2 = typename MarlinScalarType::scalar_t2; + + using FragA = typename MarlinScalarType::FragA; + using FragB = typename MarlinScalarType::FragB; + using FragC = typename MarlinScalarType::FragC; + using FragS = typename MarlinScalarType::FragS; + using FragZP = typename MarlinScalarType::FragZP; + + extern __shared__ int4 sh[]; + static constexpr auto a_type = vllm::ScalarType::from_id(a_type_id); + static constexpr auto b_type = vllm::ScalarType::from_id(b_type_id); + static constexpr auto c_type = vllm::ScalarType::from_id(c_type_id); + static constexpr auto s_type = vllm::ScalarType::from_id(s_type_id); + if constexpr (b_type == vllm::kFE2M1f) { + static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 || + s_type == vllm::kFE8M0fnu && group_blocks == 2); + } else if constexpr (b_type == vllm::kFE4M3fn && s_type == vllm::kFE8M0fnu) { + static_assert(group_blocks == 2); + } else if constexpr (std::is_same::value) { + static_assert(s_type == vllm::kBFloat16); + } else if constexpr (std::is_same::value) { + static_assert(s_type == vllm::kFloat16); + } + + constexpr bool is_a_8bit = a_type.size_bits() == 8; + if constexpr (!is_a_8bit) { + static_assert(std::is_same::value); + } + constexpr bool has_zp = b_type == vllm::kU4 || b_type == vllm::kU8; + constexpr bool is_int_type = b_type == vllm::kU4 || b_type == vllm::kU8 || + b_type == vllm::kS4 || b_type == vllm::kS8 || + b_type == vllm::kU4B8 || b_type == vllm::kU8B128; + constexpr bool is_8bit_scale = s_type.size_bits() == 8; + // see comments of dequant.h for more details + constexpr bool dequant_skip_flop = + is_a_8bit || (b_type == vllm::kFE4M3fn && !(s_type == vllm::kFE8M0fnu)) || + b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn || + has_zp && !is_zp_float && !std::is_same::value || + has_zp && !is_zp_float && !(b_type == vllm::kU8); + + float global_scale_f32 = 1.0f; + + constexpr bool has_act_order = group_blocks == 0; + + constexpr int pack_factor = 32 / b_type.size_bits(); + static_assert(thread_m_blocks == 1 || !m_block_size_8); + const int group_size = + (!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups; + const int scales_expert_stride = + prob_n * prob_k / group_size / (is_8bit_scale ? 16 : 8); + const int zp_expert_stride = + is_zp_float ? prob_n * prob_k / group_size / 8 + : prob_n * prob_k / group_size / (pack_factor * 4); + const int b_bias_expert_stride = prob_n / 8; + + // parallel: num valid moe blocks + int parallel = num_tokens_past_padded / moe_block_size; + + int k_tiles = prob_k / 16 / thread_k_blocks; + int n_tiles = prob_n / 16 / thread_n_blocks; + + int global_mn_tiles = parallel * n_tiles; + int part2_mn_tiles = global_mn_tiles; + int part1_mn_iters = 0; + bool in_part2 = false; + + // we use DP + two-tile SK here + // part1: DP + // part2: two-tile SK + // see https://github.com/vllm-project/vllm/pull/24722 for more details + if (global_mn_tiles > gridDim.x) { + part2_mn_tiles = global_mn_tiles % gridDim.x; + if (part2_mn_tiles * 3 <= gridDim.x) part2_mn_tiles += gridDim.x; + part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x; + } + + int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x); + + if constexpr (!has_act_order && group_blocks != -1) { + if (group_blocks >= thread_k_blocks) { + // Ensure that the number of tiles in each stripe is a multiple of the + // groupsize; this avoids an annoying special case where a stripe starts + // in the middle of group. + iters = (group_blocks / thread_k_blocks) * + div_ceil(iters, (group_blocks / thread_k_blocks)); + } + } + + int slice_row = 0; + int slice_col_par = blockIdx.x; + int slice_col; + int slice_iters = + k_tiles; // number of threadblock tiles in the current slice + // total number of active threadblocks in the current slice + int slice_count = 1; + // index of threadblock in current slice; numbered bottom to top + int slice_idx = 0; + + int par_id = 0; + int block_id = -1; + int64_t expert_id = 0; // use int64 to avoid computation result overflow + int old_expert_id = 0; + int64_t B_expert_off = 0; + + float* sh_a_s = reinterpret_cast(sh); + int4* sh_block_sorted_ids_int4 = sh + (is_a_8bit ? (4 * thread_m_blocks) : 0); + int4* sh_rd_block_sorted_ids_int4 = + sh_block_sorted_ids_int4 + moe_block_size / 4; + int4* sh_block_topk_weights_int4 = + sh_rd_block_sorted_ids_int4 + moe_block_size / 4; + // sh_block_topk_weights_int4 only need (moe_block_size / 4); + // but we pad to align to 256 bytes + int4* sh_new = sh_block_topk_weights_int4 + moe_block_size / 2; + int32_t* sh_block_sorted_ids = + reinterpret_cast(sh_block_sorted_ids_int4); + int32_t* sh_rd_block_sorted_ids = + reinterpret_cast(sh_rd_block_sorted_ids_int4); + c_scalar_t2* sh_block_topk_weights = + reinterpret_cast(sh_block_topk_weights_int4); + + int32_t block_num_valid_tokens = 0; + int32_t locks_off = 0; + + // We can easily implement parallel problem execution by just remapping + // indices and advancing global pointers + if (part2_mn_tiles >= gridDim.x) { + // when part2_mn_tiles >= sms + // then there are at most $sms$ conflict tile blocks + locks_off = blockIdx.x; + } else { + locks_off = (iters * blockIdx.x) / k_tiles - 1; + } + + int prob_m_top_k = prob_m * top_k; + // read moe block data given block_id + // block_sorted_ids / block_num_valid_tokens / block_topk_weights + auto read_moe_block_data = [&](int block_id) { + block_num_valid_tokens = moe_block_size; + + cp_async4_pred(sh_block_sorted_ids_int4 + threadIdx.x, + reinterpret_cast(sorted_token_ids_ptr) + + (block_id * moe_block_size / 4 + threadIdx.x), + threadIdx.x < moe_block_size / 4); + + cp_async_fence(); + cp_async_wait<0>(); + + __syncthreads(); + + if (threadIdx.x >= threads - 32) { + constexpr int size_per_thread = div_ceil(moe_block_size, 32); + int lane_id = threadIdx.x - (threads - 32); + + int local_count = 0; + #pragma unroll + for (int i = 0; i < size_per_thread; i++) { + int j = lane_id * size_per_thread + i; + if (j < moe_block_size) { + int idx = sh_block_sorted_ids[j]; + if (idx < prob_m_top_k) local_count++; + } + } + + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 + + if constexpr (moe_block_size >= 16) + local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 16); + if constexpr (moe_block_size >= 8) + local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 8); + if constexpr (moe_block_size >= 4) + local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 4); + if constexpr (moe_block_size >= 2) + local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 2); + + local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 1); + block_num_valid_tokens = local_count; + #else + block_num_valid_tokens = __reduce_add_sync(0xffffffff, local_count); + #endif + + if (lane_id == 0) + reinterpret_cast(sh_new)[0] = block_num_valid_tokens; + } + + if (threadIdx.x < moe_block_size) { + int idx = sh_block_sorted_ids[threadIdx.x]; + sh_rd_block_sorted_ids[threadIdx.x] = idx / top_k; + + if (mul_topk_weights) { + idx = idx < prob_m_top_k ? idx : 0; + float topk_weight_tmp = topk_weights_ptr[idx]; + if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) { + topk_weight_tmp *= global_scale_f32; + } + c_scalar_t2 topk_weight_val = + Cdtype::num2num2(Cdtype::float2num(topk_weight_tmp)); + sh_block_topk_weights[threadIdx.x] = topk_weight_val; + } + } + + __syncthreads(); + + block_num_valid_tokens = reinterpret_cast(sh_new)[0]; + __syncthreads(); + }; + + // when move to next moe block, find the next block_id and expert_id + // and then read moe block data + auto update_next_moe_block_data = [&]() { + if (par_id >= parallel) return; + + old_expert_id = expert_id; + block_id = par_id; + expert_id = expert_ids_ptr[block_id]; + + if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) { + global_scale_f32 = global_scale_ptr[expert_id]; + } + + B_expert_off = expert_id * prob_n * prob_k / (pack_factor * 4); + scales_ptr += (expert_id - old_expert_id) * scales_expert_stride; + if constexpr (has_zp) { + zp_ptr += (expert_id - old_expert_id) * zp_expert_stride; + } + if constexpr (has_act_order) { + g_idx += (expert_id - old_expert_id) * prob_k; + } + if (has_bias) { + b_bias_ptr += (expert_id - old_expert_id) * b_bias_expert_stride; + } + + read_moe_block_data(block_id); + }; + + // Compute all information about the current slice which is required for + // synchronization. + bool first_init = true; + auto init_part2_slice = [&]() { + slice_iters = + iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); + if (slice_iters < 0 || slice_col_par >= part2_mn_tiles) slice_iters = 0; + if (slice_iters == 0) return; + if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row; + slice_count = 1; + slice_idx = 0; + int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); + if (col_first <= k_tiles * (slice_col_par + 1)) { + int col_off = col_first - k_tiles * slice_col_par; + slice_count = div_ceil(k_tiles - col_off, iters); + if (col_off > 0) slice_count++; + int delta_first = iters * blockIdx.x - col_first; + if (delta_first < 0 || (col_off == 0 && delta_first == 0)) + slice_idx = slice_count - 1; + else { + slice_idx = slice_count - 1 - delta_first / iters; + if (col_off > 0) slice_idx--; + } + } + if (part2_mn_tiles >= gridDim.x) { + if (slice_count > 1 && slice_idx == slice_count - 1) { + locks_off++; + } + } else { + locks_off++; + } + + if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) { + constexpr int threads_per_m = 16 * thread_n_blocks / 8; + int m_per_thread = + div_ceil(block_num_valid_tokens, threads / threads_per_m); + for (int i = 0; i < m_per_thread; i++) { + int row = threads / threads_per_m * i + threadIdx.x / threads_per_m; + if (row < block_num_valid_tokens) { + int64_t sorted_row = sh_block_sorted_ids[row]; + int col = slice_col * 16 * thread_n_blocks / 8 + + threadIdx.x % threads_per_m; + C[sorted_row * prob_n / 8 + col] = {0, 0, 0, 0}; + } + } + // After write zero to output, write a negative value to lock. + // Every SM that processes the same slice would wait for + // the negative value, and then atomicAdd 1 to it. + // After all SMs are processed, the lock value would back to 0 again. + __syncthreads(); + if (threadIdx.x == 0) locks[locks_off] = 1 - slice_count; + } + + if (slice_col == n_tiles) { + slice_col = 0; + par_id++; + update_next_moe_block_data(); + } + if (is_a_8bit && (first_init || slice_col == 0)) { + __syncthreads(); + cp_async1_ca_pred(&sh_a_s[threadIdx.x], + &a_scales_ptr[sh_rd_block_sorted_ids[threadIdx.x]], + threadIdx.x < block_num_valid_tokens); + } + }; + + auto init_part1_slice = [&]() { + if (part1_mn_iters) { + part1_mn_iters--; + par_id = slice_col_par / n_tiles; + slice_col = slice_col_par % n_tiles; + slice_iters = k_tiles; + update_next_moe_block_data(); + if (is_a_8bit) { + __syncthreads(); + cp_async1_ca_pred(&sh_a_s[threadIdx.x], + &a_scales_ptr[sh_rd_block_sorted_ids[threadIdx.x]], + threadIdx.x < block_num_valid_tokens); + } + } + }; + + auto init_slice = [&]() { + if (!in_part2 && !part1_mn_iters) { + in_part2 = true; + slice_col_par = (iters * blockIdx.x) / k_tiles; + slice_row = (iters * blockIdx.x) % k_tiles; + slice_col = (slice_col_par + global_mn_tiles - part2_mn_tiles) % n_tiles; + par_id = (slice_col_par + global_mn_tiles - part2_mn_tiles) / n_tiles; + update_next_moe_block_data(); + } + if (!in_part2) { + init_part1_slice(); + } else { + init_part2_slice(); + first_init = false; + } + }; + + init_slice(); + + // A sizes/strides + + // stride of the A matrix in global memory + int a_gl_stride = prob_k / (is_a_8bit ? 16 : 8); + // stride of an A matrix tile in shared memory + constexpr int a_sh_stride = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); + // delta between subsequent A tiles in global memory + constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); + // between subsequent accesses within a tile + int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); + // between shared memory writes + constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); + // within a shared memory tile + constexpr int a_sh_rd_delta_i = a_sh_stride * 16; + // overall size of a tile + constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks); + // number of shared write iterations for a tile + constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); + + // B sizes/strides + int b_gl_stride = 16 * prob_n / (pack_factor * (is_a_8bit ? 2 : 4)); + constexpr int b_sh_stride = + ((thread_n_blocks * 16) * 16 / pack_factor) / (is_a_8bit ? 2 : 4); + constexpr int b_thread_vecs = b_type.size_bits() == 4 ? 1 : 2; + constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; + + int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); + constexpr int b_sh_wr_delta = threads * b_thread_vecs; + constexpr int b_sh_stage = + b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); + constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; + + // Scale sizes/strides without act_order + int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8); + constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8); + constexpr int s_tb_groups = + !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks + ? thread_k_blocks / group_blocks + : 1; + constexpr int s_sh_stage = s_tb_groups * s_sh_stride; + int s_gl_rd_delta = s_gl_stride; + + // Scale size/strides with act_order + constexpr int tb_k = 16 * thread_k_blocks; + constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; + // constexpr int act_s_row_stride = 1; + // int act_s_col_stride = act_s_row_stride * num_groups; + constexpr int act_s_max_num_groups = 32; + int act_s_col_stride = 1; + int act_s_col_warp_stride = act_s_col_stride * 8; + + constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4); + int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; + + // Zero-points sizes/strides + int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4; + constexpr int zp_sh_stride = is_zp_float + ? 16 * thread_n_blocks / 8 + : ((16 * thread_n_blocks) / pack_factor) / 4; + constexpr int zp_tb_groups = s_tb_groups; + constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0; + int zp_gl_rd_delta = zp_gl_stride; + + // Global A read index of current thread. + int a_gl_rd_row = threadIdx.x / a_gl_rd_delta_o; + int a_gl_rd_col = a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o; + // Shared write index of current thread. + int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + + (threadIdx.x % a_gl_rd_delta_o); + // Shared read index. + int a_sh_rd = + a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) + + (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1)); + a_sh_rd += 2 * ((threadIdx.x / 32) / tb_n_warps) * b_sh_wr_iters; + + int b_gl_rd; + if (threads <= b_sh_stride) { + b_gl_rd = threadIdx.x; + } else { + b_gl_rd = + b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); + } + + b_gl_rd += B_expert_off + b_sh_stride * slice_col; + b_gl_rd += b_gl_rd_delta_o * slice_row; + auto b_sh_rd = threadIdx.x * b_thread_vecs; + b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1)); + + // For act_order + int slice_k_start = tb_k * slice_row; + int slice_k_finish = slice_k_start + tb_k * slice_iters; + int slice_k_start_shared_fetch = slice_k_start; + int slice_n_offset = act_s_col_tb_stride * slice_col; + + // No act_order + int s_gl_rd; + if constexpr (!has_act_order) { + if constexpr (group_blocks == -1) { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + } else if constexpr (group_blocks >= thread_k_blocks) { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + s_sh_stride * slice_col + threadIdx.x; + } else { + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / s_sh_stride) + + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; + } + } + auto s_sh_wr = threadIdx.x; + bool s_sh_wr_pred = threadIdx.x < s_sh_stage; + + // Zero-points + int zp_gl_rd; + if constexpr (has_zp) { + if constexpr (group_blocks == -1) { + zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; + } else if constexpr (group_blocks >= thread_k_blocks) { + zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + zp_sh_stride * slice_col + threadIdx.x; + } else { + zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / zp_sh_stride) + + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; + } + } + auto zp_sh_wr = threadIdx.x; + bool zp_sh_wr_pred = zp_sh_stage > 0 && threadIdx.x < zp_sh_stage; + + // We use a different scale layout for grouped and column-wise quantization as + // we scale a `half2` tile in column-major layout in the former and in + // row-major in the latter case. + int s_sh_rd; + if constexpr (is_a_8bit) { + s_sh_rd = 4 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 4); + } else if constexpr (group_blocks != -1) + s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4; + else if constexpr (group_blocks == -1 && + (m_block_size_8 || (has_zp && !dequant_skip_flop))) + s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; + else + s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4; + + int bias_sh_rd; + if constexpr (m_block_size_8) { + bias_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; + } else { + bias_sh_rd = (is_a_8bit ? 4 : 8) * ((threadIdx.x / 32) % tb_n_warps) + + (threadIdx.x % 32) % 4; + } + + int bias_sh_wr = threadIdx.x; + int bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; + + // Zero-points have the same read layout as the scales + // (without column-wise case) + constexpr int num_col_threads = 8; + constexpr int num_row_threads = 4; + constexpr int num_ints_per_thread = 8 / pack_factor; + int zp_sh_rd; + if constexpr (has_zp) { + if constexpr (is_zp_float) { + if constexpr (group_blocks != -1) { + zp_sh_rd = + 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4; + } + } else if (is_a_8bit) { + zp_sh_rd = num_ints_per_thread * num_col_threads * + ((threadIdx.x / 32) % tb_n_warps / 2) + + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); + } else { + zp_sh_rd = num_ints_per_thread * num_col_threads * + ((threadIdx.x / 32) % tb_n_warps) + + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); + } + } + + // To ensure that writing and reading A tiles to/from shared memory, the + // latter in fragment format, is fully bank conflict free, we need to use a + // rather fancy XOR-based layout. The key here is that neither reads nor + // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the + // same shared memory banks. Further, it seems (based on NSight-Compute) that + // each warp must also write a consecutive memory segment? + auto transform_a = [&](int i) { + int row = i / a_gl_rd_delta_o; + return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8); + }; + // Since the computation of this remapping is non-trivial and, due to our main + // loop unrolls, all shared memory accesses are static, we simply precompute + // both transformed reads and writes. + int a_sh_wr_trans[a_sh_wr_iters]; + #pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) + a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); + int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; + #pragma unroll + for (int i = 0; i < b_sh_wr_iters; i++) { + #pragma unroll + for (int j = 0; j < thread_m_blocks; j++) + a_sh_rd_trans[i][j] = transform_a(2 * i + a_sh_rd_delta_i * j + a_sh_rd); + } + + // Since B-accesses have non-constant stride they have to be computed at + // runtime; we break dependencies between subsequent accesses with a tile by + // maintining multiple pointers (we have enough registers), a tiny + // optimization. + + // Shared memory storage for global fetch pipelines. + constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks; + constexpr int sh_b_size = stages * b_sh_stage; + int4* sh_b = sh_new; + int4* sh_red = sh_new; + + constexpr int sh_size_b_red_min = + (sh_red_size < sh_b_size ? sh_red_size : sh_b_size); + constexpr int sh_size_b_red_max = + (sh_red_size > sh_b_size ? sh_red_size : sh_b_size); + constexpr int sh_bias_size = (thread_n_blocks * 16 / 8); + constexpr int sh_b_red_bias_size = + sh_size_b_red_max > (sh_size_b_red_min + sh_bias_size) + ? sh_size_b_red_max + : (sh_size_b_red_min + sh_bias_size); + + int4* sh_bias = sh_new + sh_size_b_red_min; + int4* sh_g_idx = sh_new + sh_b_red_bias_size; + int4* sh_zp = sh_g_idx + (stages * g_idx_stage); + constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride) + : (stages * s_sh_stage); + int4* sh_s = sh_zp + (stages * zp_sh_stage); + int4* sh_a = sh_s + sh_s_size; + + // Register storage for double buffer of shared memory reads. + FragA frag_a[2][thread_m_blocks]; + I4 frag_b_quant[2][b_thread_vecs]; + FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2]; + FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2]; + FragS frag_s[2][4]; // No act-order + FragS frag_bias[2][4]; + FragS act_frag_s[2][4][4]; // For act-order + int frag_qzp[2][num_ints_per_thread]; // Zero-points + FragZP frag_zp; // Zero-points in fp16 + FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ + + if constexpr (is_a_8bit && group_blocks != -1) { + #pragma unroll + for (int j = 0; j < 2; j++) { + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int g = 0; g < 4; g++) { + frag_c_tmp[i][j][0][g] = 0.0f; + } + + #pragma unroll + for (int g = 0; g < 4; g++) { + frag_c_tmp[i][j][1][g] = 0.0f; + } + } + } + } + + // Zero accumulators. + auto zero_accums = [&]() { + #pragma unroll + for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) + reinterpret_cast(frag_c)[i] = 0; + }; + + int sh_first_group_id = -1; + int sh_num_groups = -1; + + auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id, + int last_group_id) { + sh_first_group_id = first_group_id; + sh_num_groups = last_group_id - first_group_id + 1; + + if (sh_num_groups > act_s_max_num_groups) { + sh_num_groups = act_s_max_num_groups; + } + + if (sh_first_group_id + sh_num_groups > num_groups) { + sh_num_groups = num_groups - sh_first_group_id; + } + + int row_offset = first_group_id * s_gl_stride; + + if (is_async) { + for (int i = 0; i < sh_num_groups; i++) { + if (threadIdx.x < s_sh_stride) { + cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], + &scales_ptr[row_offset + (i * s_gl_stride) + + slice_n_offset + threadIdx.x]); + } + } + } else { + for (int i = 0; i < sh_num_groups; i++) { + if (threadIdx.x < s_sh_stride) { + sh_s[(i * s_sh_stride) + threadIdx.x] = + scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + + threadIdx.x]; + } + } + } + }; + // Asynchronously fetch the next A, B and s tile from global to the next + // shared memory pipeline location. + auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { + if (pred) { + int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe; + #pragma unroll + for (int i = 0; i < a_sh_wr_iters; i++) { + int row = a_gl_rd_delta_i / a_gl_stride * i + a_gl_rd_row; + int64_t sorted_row = 0; + if (!m_block_size_8 || row < 8) + sorted_row = sh_rd_block_sorted_ids[row]; + int64_t true_idx = + sorted_row * a_gl_stride + a_gl_rd_col + a_gl_rd_delta_o * a_off; + cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], &A[true_idx], + row < block_num_valid_tokens); + } + + int4* sh_b_stage = sh_b + b_sh_stage * pipe; + #pragma unroll + for (int i = 0; i < (b_sh_wr_iters * b_thread_vecs); i++) { + constexpr int count = div_ceil(b_sh_stride, threads); + int b_gl_idx = + b_gl_rd + (i % count) * threads + + b_gl_stride * (i / count) * div_ceil(threads, b_sh_stride); + + cp_async4(&sh_b_stage[threads * i + threadIdx.x], &B[b_gl_idx]); + } + + b_gl_rd += b_gl_rd_delta_o; + + if constexpr (has_act_order) { + // Fetch g_idx thread-block portion + int full_pipe = a_off; + int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; + if (cur_k < prob_k && cur_k < slice_k_finish) { + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + + int4 const* cur_g_idx_stage_ptr = + reinterpret_cast(&g_idx[cur_k]); + + if (threadIdx.x < g_idx_stage) { + cp_async4_pred(&sh_g_idx_stage[threadIdx.x], + &cur_g_idx_stage_ptr[threadIdx.x]); + } + } + } else { + if constexpr (group_blocks != -1) { + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + // Only fetch scales if this tile starts a new group + if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { + if (s_sh_wr_pred) { + cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); + } + s_gl_rd += s_gl_rd_delta * s_tb_groups; + } + } + + if constexpr (has_zp && group_blocks != -1) { + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; + + // Only fetch zero points if this tile starts a new group + if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { + if (zp_sh_wr_pred) { + cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); + } + zp_gl_rd += zp_gl_rd_delta * zp_tb_groups; + } + } + } + } + // Insert a fence even when we are winding down the pipeline to ensure that + // waiting is also correct at this point. + cp_async_fence(); + }; + + auto fetch_col_zp_to_shared = [&]() { + if (zp_sh_wr_pred) { + cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]); + } + }; + + auto fetch_col_scale_to_shared = [&]() { + if (s_sh_wr_pred) { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + }; + + // Wait until the next thread tile has been loaded to shared memory. + auto wait_for_stage = [&]() { + // We only have `stages - 2` active fetches since we are double buffering + // and can only issue the next fetch when it is guaranteed that the previous + // shared memory load is fully complete (as it may otherwise be + // overwritten). + cp_async_wait(); + __syncthreads(); + }; + + // Load the next sub-tile from the current location in the shared memory pipe + // into the current register buffer. + auto fetch_to_registers = [&](int k, int pipe) { + int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe; + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) + ldsm( + frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); + int4* sh_b_stage = sh_b + b_sh_stage * pipe; + + #pragma unroll + for (int i = 0; i < b_thread_vecs; i++) { + frag_b_quant[k % 2][i] = *reinterpret_cast( + &sh_b_stage[b_sh_stride * (k % b_sh_wr_iters) + b_sh_rd + i]); + } + }; + + bool is_same_group[stages]; + int same_group_id[stages]; + + auto init_same_group = [&](int pipe) { + if constexpr (!has_act_order) { + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + int group_id_1 = sh_g_idx_int_ptr[0]; + int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; + + is_same_group[pipe] = group_id_1 == group_id_2; + same_group_id[pipe] = group_id_1; + }; + + auto fetch_scales_to_registers = [&](int k, int full_pipe) { + int pipe = full_pipe % stages; + using IT1 = typename std::conditional_t; + using IT0 = typename std::conditional_t; + constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1); + + if constexpr (!has_act_order) { + // No act-order case + if constexpr (group_blocks == -1) { + // load only when starting a new slice + if (k == 0 && full_pipe == 0 && dequant_skip_flop) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + } + } else if constexpr (group_blocks != -1) { + if constexpr (group_blocks >= thread_k_blocks) { + constexpr int g = group_blocks / thread_k_blocks; + if (pipe % g == 0) { + if (k % b_sh_wr_iters == 0) { + int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g)); + reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; + } else { + reinterpret_cast(&frag_s[1])[0] = + reinterpret_cast(&frag_s[0])[0]; + } + } + } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) { + auto warp_id = threadIdx.x / 32; + int warp_row = warp_id / tb_n_warps; + + int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; + int cur_group_id = k_blocks / group_blocks2; + + int4* sh_s_stage = sh_s + s_sh_stage * pipe; + + if constexpr (!is_8bit_scale) { + reinterpret_cast(&frag_s[k % 2])[0] = + sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; + } else { + reinterpret_cast(&frag_s[k % 2])[0] = + reinterpret_cast( + sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; + } + } else if (group_blocks >= b_sh_wr_iters) { + if constexpr (!is_8bit_scale) { + reinterpret_cast(&frag_s[1])[0] = + reinterpret_cast(&frag_s[0])[0]; + } else { + reinterpret_cast(&frag_s[1])[0] = + reinterpret_cast(&frag_s[0])[0]; + } + } + } + + return; + } + + // Act-order case + + // Determine K of the "current" thread-block + int cur_k = slice_k_start + tb_k * full_pipe; + if (cur_k >= prob_k || cur_k >= slice_k_finish) { + return; + } + + // Reset (to current thread-block) since we read g_idx portion from the + // shared memory + cur_k = 0; + + // Progress to current iteration + cur_k += k % b_sh_wr_iters; + + // Determine "position" inside the thread-block (based on warp and + // thread-id) + auto warp_id = threadIdx.x / 32; + int warp_row = warp_id / tb_n_warps; + int warp_col = warp_id % tb_n_warps; + + cur_k += warp_row * 16 * b_sh_wr_iters; + + auto th_id = threadIdx.x % 32; + cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix + + int s_col_shift = + /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + + (th_id / 4) * act_s_col_stride; + + if (is_same_group[pipe]) { + if (k % 2 == 0) { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = + sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + + s_col_shift]; + } else { + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = + *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); + } + + for (int i = 1; i < 4; i++) { + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = + *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); + } + return; + } + + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; + int* sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); + + constexpr int k_frag_offsets[4] = {0, 1, 8, + 9}; // Tensor core offsets per thread + + #pragma unroll + for (int i = 0; i < 4; i++) { + int actual_k = cur_k + k_frag_offsets[i]; + + int group_id = sh_g_idx_int_ptr[actual_k]; + int rel_group_id = group_id - sh_first_group_id; + + *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = + sh_s[rel_group_id * s_sh_stride + s_col_shift]; + } + }; + + auto fetch_zp_to_registers = [&](int k, int full_pipe) { + // This code does not handle group_blocks == 0, + // which signifies act_order. + // has_zp implies AWQ, which doesn't have act_order, + static_assert(!has_zp || group_blocks != 0); + + if constexpr (has_zp && !is_zp_float) { + int pipe = full_pipe % stages; + + if constexpr (group_blocks == -1) { + // load only when starting a new slice + if (k == 0 && full_pipe == 0 || is_a_8bit) { + #pragma unroll + for (int i = 0; i < num_ints_per_thread; i++) { + frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp))[zp_sh_rd + i]; + } + } + } else if constexpr (group_blocks >= thread_k_blocks) { + constexpr int g = group_blocks / thread_k_blocks; + if (pipe % g == 0 && k % b_sh_wr_iters == 0 || is_a_8bit) { + int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); + #pragma unroll + for (int i = 0; i < num_ints_per_thread; i++) { + frag_qzp[k % 2][i] = + (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; + } + } + } else { + auto warp_id = threadIdx.x / 32; + + int warp_row = warp_id / tb_n_warps; + + int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; + int cur_group_id = k_blocks / div_ceil(group_blocks, is_a_8bit ? 2 : 1); + + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; + + sh_zp_stage += cur_group_id * zp_sh_stride; + + #pragma unroll + for (int i = 0; i < num_ints_per_thread; i++) { + frag_qzp[k % 2][i] = + (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; + } + } + } + + else if constexpr (has_zp && is_zp_float) { + int pipe = full_pipe % stages; + + if constexpr (group_blocks != -1) { + if constexpr (group_blocks >= thread_k_blocks) { + constexpr int g = group_blocks / thread_k_blocks; + if (pipe % g == 0 && k % b_sh_wr_iters == 0) { + int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); + reinterpret_cast(&frag_zpf[k % 2])[0] = + sh_zp_stage[zp_sh_rd]; + } + } else if (group_blocks < b_sh_wr_iters || k % b_sh_wr_iters == 0) { + auto warp_id = threadIdx.x / 32; + + int warp_row = warp_id / tb_n_warps; + int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; + int cur_group_id = k_blocks / group_blocks; + + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; + + reinterpret_cast(&frag_zpf[k % 2])[0] = + sh_zp_stage[zp_sh_rd + cur_group_id * zp_sh_stride]; + } + } + } + }; + + auto dequant_data = [&](int q, scalar_32bit_t* frag_b_ptr, int zp = 0) { + if constexpr (a_type.size_bits() != b_type.size_bits()) { + if constexpr (is_a_8bit && has_zp) { + sub_zp_and_dequant( + q, frag_b_ptr, zp); + } else { + dequant(q, frag_b_ptr); + } + } + }; + + // Execute the actual tensor core matmul of a sub-tile. + bool is_first_matmul_in_slice = true; + auto matmul = [&](int k, int pipe) { + if (is_a_8bit) return; + int k2 = k % 2; + constexpr int g = + group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1; + const bool is_new_zp = + (group_blocks == 0) || + ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) && + (pipe % g == 0) || + (group_blocks == -1 && is_first_matmul_in_slice); + if constexpr (has_zp && !is_zp_float) { + if (is_new_zp) { + if constexpr (group_blocks == -1) is_first_matmul_in_slice = false; + int zp_quant_0, zp_quant_1; + + if constexpr (b_type.size_bits() == 4) { + zp_quant_0 = frag_qzp[k2][0]; + zp_quant_1 = zp_quant_0 >> 8; + } else { + static_assert(b_type.size_bits() == 8); + zp_quant_0 = frag_qzp[k2][0]; + zp_quant_1 = frag_qzp[k2][1]; + } + + dequant_data(zp_quant_0, reinterpret_cast(&frag_zp)); + dequant_data(zp_quant_1, + reinterpret_cast(&frag_zp) + 2); + } + } + if constexpr (!dequant_skip_flop && has_zp && is_zp_float) { + if (is_new_zp) { + reinterpret_cast(&frag_zp)[0] = + reinterpret_cast(&frag_zpf[k2])[0]; + } + } + + if constexpr (s_type == vllm::kFE4M3fn || s_type == vllm::kFE8M0fnu) { + int s_quant_0 = reinterpret_cast(frag_s[k2])[0]; + int s_quant_1 = reinterpret_cast(frag_s[k2])[1]; + + dequant_fp8_scales( + s_quant_0, reinterpret_cast(&frag_s[k2])); + dequant_fp8_scales( + s_quant_1, reinterpret_cast(&frag_s[k2]) + 2); + } + + // We have the m dimension as the inner loop in order to encourage overlapping + // dequantization and matmul operations. + #pragma unroll + for (int j = 0; j < 4; j++) { + FragB frag_b0; + FragB frag_b1; + int b_quant_0, b_quant_1; + + if constexpr (b_type_id == vllm::kFE2M1f.id()) { + b_quant_1 = frag_b_quant[k2][0][j]; + b_quant_0 = b_quant_1 << 8; + } else if constexpr (b_type.size_bits() == 4) { + b_quant_0 = frag_b_quant[k2][0][j]; + b_quant_1 = b_quant_0 >> 8; + } else { + static_assert(b_type.size_bits() == 8); + int* frag_b_quant_ptr = reinterpret_cast(frag_b_quant[k2]); + b_quant_0 = frag_b_quant_ptr[j * 2 + 0]; + b_quant_1 = frag_b_quant_ptr[j * 2 + 1]; + } + + dequant_data(b_quant_0, reinterpret_cast(&frag_b0)); + dequant_data(b_quant_1, reinterpret_cast(&frag_b1)); + + if constexpr (dequant_skip_flop && has_zp && !is_zp_float && !is_a_8bit) { + sub_zp(frag_b0, frag_zp[j], 0); + sub_zp(frag_b1, frag_zp[j], 1); + } + + // Apply scale to frag_b0 + if constexpr (has_act_order && !is_a_8bit) { + static_assert(group_blocks != -1); + scale4(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j], + act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0); + scale4(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j], + act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1); + } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float && + group_blocks == -1 && !is_a_8bit) { + int idx = (threadIdx.x / 4) % 2; + scalar_t2 s2 = Adtype::nums2num2( + reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], + reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 1])[idx]); + if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2); + scale_and_sub(frag_b0, s2.x, frag_zp[j].x); + scale_and_sub(frag_b1, s2.y, frag_zp[j].y); + } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1 && + !is_a_8bit) { + if (is_new_zp) + frag_zp[j] = __hmul2(frag_zp[j], + *reinterpret_cast(&frag_s[k2][j])); + scale_and_sub(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x); + scale_and_sub(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y); + } else if constexpr (group_blocks != -1 && !is_a_8bit) { + scale(frag_b0, frag_s[k2][j], 0); + scale(frag_b1, frag_s[k2][j], 1); + } + + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + if constexpr (m_block_size_8) { + mma_trans(frag_a[k2][i], frag_b0, frag_b1, + frag_c[i][j][0]); + } else { + mma(frag_a[k2][i], frag_b0, + frag_c[i][j][0]); + mma(frag_a[k2][i], frag_b1, + frag_c[i][j][1]); + } + } + } + }; + + auto matmul_a8 = [&](int k) { + int k2 = k % 2; + #pragma unroll + for (int j = 0; j < 2; j++) { + FragB frag_b[2]; + + if (is_a_8bit && b_type.size_bits() == 4 && !has_zp) { + dequant_data(frag_b_quant[k2][0][j * 2], + reinterpret_cast(&frag_b)); + dequant_data(frag_b_quant[k2][0][j * 2 + 1], + reinterpret_cast(&frag_b) + 2); + } else if (is_a_8bit && b_type.size_bits() == 4 && has_zp) { + int off = (threadIdx.x / 32) % 2 * 2 + j; + int zp = (frag_qzp[k2][0] >> (off * 8)) & 0xF; + dequant_data(frag_b_quant[k2][0][j * 2], + reinterpret_cast(&frag_b), zp); + zp = (frag_qzp[k2][0] >> (off * 8 + 4)) & 0xF; + dequant_data(frag_b_quant[k2][0][j * 2 + 1], + reinterpret_cast(&frag_b) + 2, zp); + } else { + reinterpret_cast(&frag_b)[0] = + reinterpret_cast(&frag_b_quant[k2][j])[0]; + reinterpret_cast(&frag_b)[1] = + reinterpret_cast(&frag_b_quant[k2][j])[1]; + } + + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + mma( + frag_a[k2][i], frag_b[0], + (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]); + mma( + frag_a[k2][i], frag_b[1], + (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]); + } + + if constexpr (group_blocks != -1) { + if (group_blocks == 2 || k == 1) { + if constexpr (a_type == vllm::kS8) { + int2 s_vals[2]; + s_vals[0] = { + (int)reinterpret_cast(&frag_s[k2][j * 2][0])[0], + (int)reinterpret_cast(&frag_s[k2][j * 2][0])[1]}; + s_vals[1] = { + (int)reinterpret_cast(&frag_s[k2][j * 2 + 1][0])[0], + (int)reinterpret_cast(&frag_s[k2][j * 2 + 1][0])[1]}; + + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int g = 0; g < 4; g++) { + int scale = reinterpret_cast(&s_vals[0])[g % 2]; + *reinterpret_cast(&frag_c[i][j][0][g]) += + *reinterpret_cast(&frag_c_tmp[i][j][0][g]) * + scale; + frag_c_tmp[i][j][0][g] = 0.0f; + } + + #pragma unroll + for (int g = 0; g < 4; g++) { + int scale = reinterpret_cast(&s_vals[1])[g % 2]; + *reinterpret_cast(&frag_c[i][j][1][g]) += + *reinterpret_cast(&frag_c_tmp[i][j][1][g]) * + scale; + frag_c_tmp[i][j][1][g] = 0.0f; + } + } + } else { + float2 s_vals[2]; + if constexpr (s_type_id != vllm::kFE8M0fnu.id()) { + static_assert(a_type.size_bits() == 16 || + s_type.size_bits() == 16); + s_vals[0] = Cdtype::num22float2(frag_s[k2][j * 2][0]); + s_vals[1] = Cdtype::num22float2(frag_s[k2][j * 2 + 1][0]); + } else { + int32_t* s_vals_int = reinterpret_cast(&s_vals[0]); + int32_t s_vals_e8m0 = + *reinterpret_cast(&frag_s[k2][j][0]); + + s_vals_int[0] = (s_vals_e8m0 & 0xFF) << 23; + s_vals_int[1] = (s_vals_e8m0 & 0xFF00) << 15; + s_vals_int[2] = (s_vals_e8m0 & 0xFF0000) << 7; + s_vals_int[3] = (s_vals_e8m0 & 0xFF000000) >> 1; + } + + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int g = 0; g < 4; g++) { + float scale = reinterpret_cast(&s_vals[0])[g % 2]; + frag_c[i][j][0][g] += frag_c_tmp[i][j][0][g] * scale; + frag_c_tmp[i][j][0][g] = 0.0f; + } + + #pragma unroll + for (int g = 0; g < 4; g++) { + float scale = reinterpret_cast(&s_vals[1])[g % 2]; + frag_c[i][j][1][g] += frag_c_tmp[i][j][1][g] * scale; + frag_c_tmp[i][j][1][g] = 0.0f; + } + } + } + } + } + } + }; + + // Since we slice across the k dimension of a tile in order to increase the + // number of warps while keeping the n dimension of a tile reasonable, we have + // multiple warps that accumulate their partial sums of the same output + // location; which we have to reduce over in the end. We do in shared memory. + auto thread_block_reduce = [&]() { + constexpr int red_off = threads / b_sh_stride_threads / 2; + if (red_off >= 1) { + auto red_idx = threadIdx.x / b_sh_stride_threads; + constexpr int red_sh_stride = + b_sh_stride_threads * (is_a_8bit ? 2 : 4) * 2; + constexpr int red_sh_delta = b_sh_stride_threads; + int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + + (threadIdx.x % b_sh_stride_threads); + + // Parallel logarithmic shared memory reduction. We make sure to avoid any + // unnecessary read or write iterations, e.g., for two warps we write only + // once by warp 1 and read only once by warp 0. + + #pragma unroll + for (int m_block = 0; m_block < thread_m_blocks; m_block++) { + #pragma unroll + for (int i = red_off; i > 0; i /= 2) { + if (i <= red_idx && red_idx < 2 * i) { + #pragma unroll + for (int j = 0; j < (is_a_8bit ? 2 : 4) * 2; + j += (m_block_size_8 ? 2 : 1)) { + int red_sh_wr = + red_sh_delta * j + (red_sh_rd - red_sh_stride * i); + if (i < red_off) { + float* c_rd = reinterpret_cast( + &sh_red[red_sh_delta * j + red_sh_rd]); + float* c_wr = reinterpret_cast(&sh_red[red_sh_wr]); + #pragma unroll + for (int k = 0; k < 4; k++) + reinterpret_cast( + frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j][k] += + c_rd[k] + c_wr[k]; + } + sh_red[red_sh_wr] = reinterpret_cast( + &frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j]; + } + } + __syncthreads(); + } + if (red_idx == 0) { + #pragma unroll + for (int i = 0; i < (is_a_8bit ? 2 : 4) * 2; + i += (m_block_size_8 ? 2 : 1)) { + float* c_rd = + reinterpret_cast(&sh_red[red_sh_delta * i + red_sh_rd]); + #pragma unroll + for (int j = 0; j < 4; j++) + reinterpret_cast( + frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + i][j] += c_rd[j]; + } + } + __syncthreads(); + } + } + }; + + // Since multiple threadblocks may process parts of the same column slice, we + // finally have to globally reduce over the results. As the striped + // partitioning minimizes the number of such reductions and our outputs are + // usually rather small, we perform this reduction serially in L2 cache. + auto global_reduce_fp16 = [&](bool first = false, bool last = false) { + // We are very careful here to reduce directly in the output buffer to + // maximize L2 cache utilization in this step. To do this, we write out + // results in FP16 (but still reduce with FP32 compute). + constexpr int active_threads = 32 * tb_n_warps; + bool is_th_active = threadIdx.x < active_threads; + if (!is_th_active) { + return; + } + + int c_gl_stride = prob_n / 8 * (is_a_8bit ? 2 : 1); + int c_gl_wr_delta_o = 8 * c_gl_stride; + int c_gl_wr_delta_i = 4 * (active_threads / 32); + int c_gl_wr; + if constexpr (m_block_size_8) { + c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + 4 * (threadIdx.x / 32) + + (threadIdx.x % 32) / 8; + c_gl_wr += (2 * thread_n_blocks) * slice_col; + } else { + c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) + + 4 * (threadIdx.x / 32) + threadIdx.x % 4; + c_gl_wr += (2 * thread_n_blocks) * slice_col * (is_a_8bit ? 2 : 1); + } + constexpr int c_sh_wr_delta = active_threads; + int c_sh_wr = threadIdx.x; + + if (!first) { + + #pragma unroll + for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { + int c_idx; + if constexpr (m_block_size_8) + c_idx = c_gl_wr + i * c_gl_stride + + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i; + else + c_idx = + c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2); + if (c_idx / c_gl_stride < block_num_valid_tokens) { + int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride]; + int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride; + if constexpr (is_a_8bit) { + int2* sh_red_int2 = reinterpret_cast(sh_red); + int2* c_int2 = reinterpret_cast(C); + sh_red_int2[c_sh_wr + c_sh_wr_delta * i] = c_int2[true_idx]; + } else { + sh_red[c_sh_wr + c_sh_wr_delta * i] = C[true_idx]; + } + } + } + } + + #pragma unroll + for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { + if (!first) { + c_scalar_t* c_red_f16; + if constexpr (is_a_8bit) { + int2 tmp = + reinterpret_cast(sh_red)[c_sh_wr + i * c_sh_wr_delta]; + c_red_f16 = reinterpret_cast(&tmp); + } else { + int4 tmp = sh_red[c_sh_wr + i * c_sh_wr_delta]; + c_red_f16 = reinterpret_cast(&tmp); + } + #pragma unroll + for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) { + int delta = 0; + if constexpr (m_block_size_8) { + delta = j % 2 == 1 ? -2 : 0; + } + reinterpret_cast( + &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + + delta] += Cdtype::num2float(c_red_f16[j]); + } + } + if (!last) { + c_scalar_t c_f16[is_a_8bit ? 4 : 8]; + #pragma unroll + for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) { + int delta = 0; + if constexpr (m_block_size_8) { + delta = j % 2 == 1 ? -2 : 0; + } + c_f16[j] = Cdtype::float2num(reinterpret_cast( + &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + + delta]); + } + + int c_idx; + if constexpr (m_block_size_8) + c_idx = c_gl_wr + i * c_gl_stride + + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i; + else + c_idx = + c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2); + if (c_idx / c_gl_stride < block_num_valid_tokens) { + int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride]; + int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride; + if constexpr (is_a_8bit) { + int2* c_int2 = reinterpret_cast(C); + c_int2[true_idx] = *reinterpret_cast(c_f16); + } else { + C[true_idx] = *reinterpret_cast(c_f16); + } + } + } + } + }; + + // Globally reduce over threadblocks that compute the same column block. + // We use a tmp C buffer to reduce in full fp32 precision. + auto global_reduce_fp32 = [&](bool first = false, bool last = false) { + constexpr int tb_m = thread_m_blocks * 16; + constexpr int tb_n = thread_n_blocks * 16; + + constexpr int c_size = tb_m * tb_n * sizeof(float) / 16; + + constexpr int active_threads = 32 * tb_n_warps; + bool is_th_active = threadIdx.x < active_threads; + + constexpr int num_floats = thread_m_blocks * (is_a_8bit ? 2 : 4) * 2 * 4; + constexpr int th_size = num_floats * sizeof(float) / 16; + + int c_cur_offset = locks_off * c_size; + + if (!is_th_active) { + return; + } + + if (!first) { + float* frag_c_ptr = reinterpret_cast(&frag_c); + #pragma unroll + for (int k = 0; k < th_size; k++) { + if constexpr (m_block_size_8) { + if (k % 2) continue; + } else { + if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens) + continue; + } + + sh_red[threadIdx.x] = + C_tmp[c_cur_offset + active_threads * k + threadIdx.x]; + + float* sh_c_ptr = reinterpret_cast(&sh_red[threadIdx.x]); + #pragma unroll + for (int f = 0; f < 4; f++) { + frag_c_ptr[k * 4 + f] += sh_c_ptr[f]; + } + } + } + + if (!last) { + int4* frag_c_ptr = reinterpret_cast(&frag_c); + #pragma unroll + for (int k = 0; k < th_size; k++) { + if constexpr (m_block_size_8) { + if (k % 2) continue; + } else { + if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens) + continue; + } + + C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k]; + } + } + }; + + // Write out the reduce final result in the correct layout. We only actually + // reshuffle matrix fragments in this step, the reduction above is performed + // in fragment layout. + auto write_result = [&](bool last) { + int c_gl_stride = prob_n / 8; + constexpr int c_sh_stride = 2 * thread_n_blocks + 1; + int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); + constexpr int c_sh_rd_delta = + c_sh_stride * (threads / (2 * thread_n_blocks)); + + int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + + (threadIdx.x % (2 * thread_n_blocks)); + c_gl_wr += (2 * thread_n_blocks) * slice_col; + int c_sh_wr; + if constexpr (m_block_size_8) { + c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) + + (threadIdx.x % 32) / 4; + c_sh_wr += 64 * (threadIdx.x / 32); + } else { + c_sh_wr = + (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; + c_sh_wr += (is_a_8bit ? 16 : 32) * (threadIdx.x / 32); + } + + int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + + (threadIdx.x % (2 * thread_n_blocks)); + + // We first reorder in shared memory to guarantee the most efficient final + // global write patterns + auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) { + if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) { + if (!mul_topk_weights) { + c0 *= global_scale_f32; + c1 *= global_scale_f32; + } + } + + c_scalar_t2 res = + Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1)); + + // For per-column quantization we finally apply the scale here (only for + // 4-bit) + if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit && + b_type.size_bits() == 4 && + (has_zp && dequant_skip_flop || !has_zp)) { + c_scalar_t2 tmp_scale = s[0]; + if constexpr (m_block_size_8) { + tmp_scale = Cdtype::num2num2( + reinterpret_cast(&s[0])[(threadIdx.x % 8) / 4]); + } + res = __hmul2(res, tmp_scale); + } + + if (has_bias && last) { + c_scalar_t2 tmp_bias = b_bias[0]; + if constexpr (m_block_size_8) { + tmp_bias = Cdtype::num2num2( + reinterpret_cast(&b_bias[0])[(threadIdx.x % 8) / 4]); + } + res = __hadd2(res, tmp_bias); + } + + if constexpr (m_block_size_8) { + ((c_scalar_t*)sh_red)[idx] = res.x; + ((c_scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y; + } else { + ((c_scalar_t2*)sh_red)[idx] = res; + } + }; + + if (threadIdx.x / 32 < tb_n_warps) { + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int j = 0; j < (is_a_8bit ? 2 : 4); j++) { + if constexpr (m_block_size_8) { + int wr = c_sh_wr + 16 * j; + write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1], + frag_s[j / 2][2 * (j % 2) + 0], + frag_bias[j / 2][2 * (j % 2) + 0]); + write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3], + frag_s[j / 2][2 * (j % 2) + 1], + frag_bias[j / 2][2 * (j % 2) + 1]); + } else { + int wr = c_sh_wr + 8 * j; + write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], + frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0], + frag_bias[j / 2][2 * (j % 2) + 0]); + write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], + frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0], + frag_bias[j / 2][2 * (j % 2) + 0]); + write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], + frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1], + frag_bias[j / 2][2 * (j % 2) + 1]); + write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], + frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1], + frag_bias[j / 2][2 * (j % 2) + 1]); + } + } + c_sh_wr += 16 * (4 * c_sh_stride); + } + } + __syncthreads(); + + #pragma unroll + for (int i = 0; + i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); + i++) { + int row = c_gl_wr / c_gl_stride; + if (row < block_num_valid_tokens) { + int64_t sorted_row = sh_block_sorted_ids[row]; + int64_t true_idx = sorted_row * c_gl_stride + c_gl_wr % c_gl_stride; + c_scalar_t2 topk_weight_score; + if (mul_topk_weights) topk_weight_score = sh_block_topk_weights[row]; + if (use_atomic_add && slice_count > 1 || mul_topk_weights) { + c_scalar_t2* C_half2 = reinterpret_cast(&C[true_idx]); + c_scalar_t2* sh_red_half2 = + reinterpret_cast(&sh_red[c_sh_rd]); + if (mul_topk_weights) { + #pragma unroll + for (int a = 0; a < 4; a++) { + sh_red_half2[a] = __hmul2(sh_red_half2[a], topk_weight_score); + } + } + + if (use_atomic_add && slice_count > 1) { + #pragma unroll + for (int a = 0; a < 4; a++) { + atomicAdd(&C_half2[a], sh_red_half2[a]); + } + } else { + C[true_idx] = *reinterpret_cast(sh_red_half2); + } + } else { + C[true_idx] = sh_red[c_sh_rd]; + } + c_gl_wr += c_gl_wr_delta; + c_sh_rd += c_sh_rd_delta; + } + } + __syncthreads(); + }; + + // Start global fetch and register load pipelines. + auto start_pipes = [&]() { + + #pragma unroll + for (int i = 0; i < stages - 1; i++) { + if (has_act_order && i == 0) { + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) { + last_g_idx = prob_k - 1; + } + fetch_act_order_scales_to_shared(true, g_idx[slice_k_start], + g_idx[last_g_idx]); + } + + if constexpr (has_zp && !is_zp_float && group_blocks == -1) { + if (i == 0) { + fetch_col_zp_to_shared(); + if constexpr (!dequant_skip_flop) { + fetch_col_scale_to_shared(); + } + } + } + fetch_to_shared(i, i, i < slice_iters); + } + + zero_accums(); + wait_for_stage(); + init_same_group(0); + fetch_to_registers(0, 0); + fetch_scales_to_registers(0, 0); + fetch_zp_to_registers(0, 0); + a_gl_rd_col += a_gl_rd_delta_o * (stages - 1); + if constexpr (has_act_order) { + slice_k_start_shared_fetch += tb_k * (stages - 1); + } + }; + if (slice_iters) { + start_pipes(); + } + + // Main loop. + while (slice_iters) { + // We unroll over both the global fetch and the register load pipeline to + // ensure all shared memory accesses are static. Note that both pipelines + // have even length meaning that the next iteration will always start at + // index 0. + + #pragma unroll + for (int pipe = 0; pipe < stages;) { + #pragma unroll + for (int k = 0; k < b_sh_wr_iters; k++) { + fetch_to_registers(k + 1, pipe % stages); + fetch_scales_to_registers(k + 1, pipe); + fetch_zp_to_registers(k + 1, pipe); + if (k == b_sh_wr_iters - 2) { + fetch_to_shared((pipe + stages - 1) % stages, pipe, + slice_iters >= stages); + pipe++; + wait_for_stage(); + init_same_group(pipe % stages); + } + + if constexpr (!is_a_8bit) { + matmul(k, pipe - (k >= b_sh_wr_iters - 2 ? 1 : 0)); + } else { + static_assert(group_blocks != 0 && group_blocks != 1); + matmul_a8(k); + } + } + slice_iters--; + if (slice_iters == 0) { + break; + } + } + + a_gl_rd_col += a_gl_rd_delta_o * stages; + + if constexpr (has_act_order) { + slice_k_start += tb_k * stages; + + if (slice_k_start < prob_k) { + slice_k_start_shared_fetch += tb_k * stages; + int first_group_id = g_idx[slice_k_start]; + int last_g_idx = slice_k_start + stages * tb_k * 2; + if (last_g_idx >= prob_k) { + last_g_idx = prob_k - 1; + } + int last_group_id = g_idx[last_g_idx]; + if (last_group_id >= sh_first_group_id + sh_num_groups) { + fetch_act_order_scales_to_shared(false, first_group_id, + last_group_id); + __syncthreads(); + } + } + } + + // Process results and, if necessary, proceed to the next column slice. + // While this pattern may not be the most readable, other ways of writing + // the loop seemed to noticeably worse performance after compilation. + if (slice_iters == 0) { + // convert fp16 accum to fp32 for reduction + if constexpr (use_fp16_accum) { + #pragma unroll + for (int i = 0; i < (thread_m_blocks * (is_a_8bit ? 2 : 4) * 2); i++) { + float* frag_c_part_float = reinterpret_cast(frag_c) + i * 4; + scalar_t* frag_c_part_half = + reinterpret_cast(frag_c_part_float); + + #pragma unroll + for (int i = 3; i >= 0; i--) { + frag_c_part_float[i] = Cdtype::num2float(frag_c_part_half[i]); + } + } + } + + if constexpr (is_a_8bit) { + float frag_a_s[2 * thread_m_blocks]; + + for (int i = 0; i < 2 * thread_m_blocks; i++) + frag_a_s[i] = sh_a_s[i * 8 + (threadIdx.x % 32) / 4]; + + #pragma unroll + for (int j = 0; j < 2; j++) { + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int g = 0; g < 4; g++) { + float c_val = frag_c[i][j][0][g]; + + if constexpr (a_type == vllm::kS8) { + c_val = __int2float_rn(*reinterpret_cast(&c_val)); + } + float s_val = frag_a_s[i * 2 + g / 2]; + frag_c[i][j][0][g] = c_val * s_val; + } + #pragma unroll + for (int g = 0; g < 4; g++) { + float c_val = frag_c[i][j][1][g]; + + if constexpr (a_type == vllm::kS8) { + c_val = __int2float_rn(*reinterpret_cast(&c_val)); + } + float s_val = frag_a_s[i * 2 + g / 2]; + frag_c[i][j][1][g] = c_val * s_val; + } + } + } + } + + cp_async_wait<0>(); + bool last = slice_idx == slice_count - 1; + // For per-column scales, we only fetch them here in the final step before + // write-out + if constexpr (!has_act_order && group_blocks == -1 && + (has_zp && dequant_skip_flop || !has_zp)) { + if (b_type.size_bits() == 8 || (last || use_atomic_add) || is_a_8bit) { + if (s_sh_wr_pred) { + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); + } + cp_async_fence(); + } + } + + thread_block_reduce(); + + if (has_bias && last) { + __syncthreads(); + cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd], + threadIdx.x < 16 * thread_n_blocks / 8); + cp_async_fence(); + } + + if constexpr (!has_act_order && group_blocks == -1 && + (has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) { + if constexpr (is_a_8bit) { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < tb_n_warps) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + } + } else if (b_type.size_bits() == 8 || (last || use_atomic_add)) { + cp_async_wait<0>(); + __syncthreads(); + if (threadIdx.x / 32 < tb_n_warps) { + reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; + reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; + if constexpr (m_block_size_8) { + int idx = (threadIdx.x / 4) % 2; + c_scalar_t2* frag_s_half2 = + reinterpret_cast(frag_s); + #pragma unroll + for (int i = 0; i < 8; i++) { + frag_s_half2[i] = Cdtype::num2num2( + reinterpret_cast(&frag_s_half2[i])[idx]); + } + } + } + } + } + + // For 8-bit channelwise, we apply the scale before the global reduction + // that converts the fp32 results to fp16 (so that we avoid possible + // overflow in fp16) + if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) { + #pragma unroll + for (int j = 0; j < 2; j++) { + float2 aa[2]; + aa[0] = Cdtype::num22float2(frag_s[0][j * 2][0]); + aa[1] = Cdtype::num22float2(frag_s[0][j * 2 + 1][0]); + + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int g = 0; g < 4; g++) { + float scale = reinterpret_cast(&aa[0])[g % 2]; + frag_c[i][j][0][g] *= scale; + } + + #pragma unroll + for (int g = 0; g < 4; g++) { + float scale = reinterpret_cast(&aa[1])[g % 2]; + frag_c[i][j][1][g] *= scale; + } + } + } + } else if (!has_act_order && group_blocks == -1 && + b_type.size_bits() == 8 && + (has_zp && dequant_skip_flop || !has_zp)) { + if (threadIdx.x / 32 < tb_n_warps) { + #pragma unroll + for (int i = 0; i < thread_m_blocks; i++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + scale_float( + reinterpret_cast(&frag_c[i][j][0][0]), + frag_s[j / 2][2 * (j % 2) + 0]); + scale_float( + reinterpret_cast(&frag_c[i][j][0][2]), + frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]); + + if constexpr (!m_block_size_8) { + scale_float( + reinterpret_cast(&frag_c[i][j][1][0]), + frag_s[j / 2][2 * (j % 2) + 1]); + scale_float( + reinterpret_cast(&frag_c[i][j][1][2]), + frag_s[j / 2][2 * (j % 2) + 1]); + } + } + } + } + } + + if (slice_count > 1 && !use_atomic_add) { + // only globally reduce if there is more than one block in a slice + barrier_acquire(&locks[locks_off], slice_idx); + if (use_fp32_reduce) { + global_reduce_fp32(slice_idx == 0, last); + } else { + global_reduce_fp16(slice_idx == 0, last); + } + barrier_release(&locks[locks_off], last); + } + + if (has_bias && last) { + cp_async_wait<0>(); + __syncthreads(); + reinterpret_cast(&frag_bias)[0] = sh_bias[bias_sh_rd]; + if constexpr (!is_a_8bit) + reinterpret_cast(&frag_bias)[1] = sh_bias[bias_sh_rd + 4]; + __syncthreads(); + } + + if (use_atomic_add && slice_count > 1 && slice_idx != 0) + wait_negative_and_add(&locks[locks_off]); + if (last || use_atomic_add) + // only the last block in a slice actually writes the result + write_result(last); + slice_row = 0; + if (!in_part2) { + slice_col_par += gridDim.x; + } else { + slice_col_par++; + slice_col++; + } + is_first_matmul_in_slice = true; + init_slice(); + + if (slice_iters) { + a_gl_rd_col = + a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o; + b_gl_rd = B_expert_off + b_gl_stride * (threadIdx.x / b_sh_stride) + + (threadIdx.x % b_sh_stride); + b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row; + + bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; + // Update slice k/n for scales loading + if constexpr (has_act_order) { + slice_k_start = tb_k * slice_row; + slice_k_finish = slice_k_start + tb_k * slice_iters; + slice_k_start_shared_fetch = slice_k_start; + slice_n_offset = act_s_col_tb_stride * slice_col; + } else { + if constexpr (group_blocks == -1) { + s_gl_rd = s_sh_stride * slice_col + threadIdx.x; + zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; + } else if constexpr (group_blocks >= thread_k_blocks) { + s_gl_rd = + s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + s_sh_stride * slice_col + threadIdx.x; + zp_gl_rd = + zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + + zp_sh_stride * slice_col + threadIdx.x; + } else { + s_gl_rd = + s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / s_sh_stride) + + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; + zp_gl_rd = + zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + + threadIdx.x / zp_sh_stride) + + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; + } + } + start_pipes(); + } + } + } +} + +} // namespace MARLIN_NAMESPACE_NAME + +#endif diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_u4b8_bfloat16.cu b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_u4b8_bfloat16.cu new file mode 100644 index 000000000..0948a68f8 --- /dev/null +++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_u4b8_bfloat16.cu @@ -0,0 +1,160 @@ +// auto generated by generate_kernels.py +// clang-format off + +#include "kernel.h" +#include "marlin_template.h" + +namespace MARLIN_NAMESPACE_NAME { + + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +} diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_float16_u4b8_float16.cu b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_float16_u4b8_float16.cu new file mode 100644 index 000000000..fdace4b09 --- /dev/null +++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_float16_u4b8_float16.cu @@ -0,0 +1,160 @@ +// auto generated by generate_kernels.py +// clang-format off + +#include "kernel.h" +#include "marlin_template.h" + +namespace MARLIN_NAMESPACE_NAME { + + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +template __global__ void Marlin( MARLIN_KERNEL_PARAMS ); + +} diff --git a/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/dequant.h b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/dequant.h new file mode 100644 index 000000000..edd97dbfc --- /dev/null +++ b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/dequant.h @@ -0,0 +1,609 @@ +/* +Fast Dequantization (Converting INT4/INT8/FP4/FP8 to FP16/BF16) + +The process of fast dequantization can be summarized as a combination +of bitwise operations and floating-point computations: + +weight =>(bit_op / bitwise operations)=> +f16_value =>(flop / floating-point computation)=> +dequantized_weight + +Since the dequantized weights typically require subtracting the zero point and +applying a scale factor, the floating-point computation step can be fused with +the zero-point subtraction and scaling operations. + +The following are the parts that need to be modified for the fused operation +of zero-point subtraction and scaling. + +## INT4 => FP16/BF16 or INT8 => FP16 + +The floating-point computation is `__hsub2` + +If has zero points: + + flop(bit_op(weight)) - flop(bit_op(zp)) + = sub(bit_op(weight), bias) - sub(bit_op(zp), bias) + = bit_op(weight) - bit_op(zp) + +so we don't need additional modification. + +If has float zero points: + + flop(bit_op(weight)) - fzp + = sub(bit_op(weight), bias) - fzp + = bit_op(weight) - (fzp + bias) + +where the `fzp + bias` can be computed at weight loading. But this +may have accuracy issue, so we should not use this in most cases. + +If has not zero points: + + scale(flop(bit_op(weight))) + = scale(sub(bit_op(weight), bias)) + = scale(bit_op(weight)) - scale(bias) + = fma(bit_op(weight), scale_factor, scale(bias)) + +where the `scale(bias)` can be cached. But this may have accuracy issue, +so we should not use this in most cases. + + +## INT8 => BF16 + +INT8 => BF16 is a special case, it use byte_perm instead of flop. +We cannot fused byte_perm with scaling. + + +## FP4/FP8 => FP16/BF16 + + scale(flop(bit_op(weight))) + = scale(mul(bit_op(weight), multiplier)) + = mul(bit_op(weight), scale_factor * multiplier) + +where `scale_factor * multiplier` can be computed at weight loading. + +*/ + +#include "marlin_dtypes.cuh" + +namespace MARLIN_NAMESPACE_NAME { + +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 750 +// Lookup-table based 3-input logical operation; explicitly used for +// dequantization as the compiler does not seem to automatically recognize it in +// all cases. +template +__device__ inline int lop3(int a, int b, int c) { + int res; + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(res) + : "r"(a), "r"(b), "r"(c), "n"(lut)); + return res; +} + +// Constructs destination register by taking bytes from 2 sources (based on +// mask) +template +__device__ inline uint32_t prmt(uint32_t a) { + uint32_t res; + asm volatile("prmt.b32 %0, %1, %2, %3;\n" + : "=r"(res) + : "r"(a), "n"(start_byte), "n"(mask)); + return res; +} + +template +__device__ inline void dequant(int q, scalar_t2* frag_b); + +// +// Efficiently dequantize 4bit values packed in an int32 value into a full +// B-fragment of 4 fp16 values. We mostly follow the strategy in the link below, +// with some small changes: +// - FP16: +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287 +// - BF16: +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385 +// +template <> +__device__ inline void dequant(int q, + half2* frag_b) { + const int MASK = 0x000f000f; + const int EX = 0x64006400; + // Guarantee that the `(a & b) | c` operations are LOP3s. + int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); + q >>= 4; + int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); + + frag_b[0] = *reinterpret_cast(&lo); + frag_b[1] = *reinterpret_cast(&hi); +} + +template <> +__device__ inline void dequant(int q, + half2* frag_b) { + const int LO = 0x000f000f; + const int HI = 0x00f000f0; + const int EX = 0x64006400; + // Guarantee that the `(a & b) | c` operations are LOP3s. + // clang-format off + int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); + int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); + // clang-format on + // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point + // directly into `SUB` and `ADD`. + const int SUB = 0x64086408; + const int MUL = 0x2c002c00; + const int ADD = 0xd480d480; + frag_b[0] = __hsub2(*reinterpret_cast(&lo), + *reinterpret_cast(&SUB)); + frag_b[1] = __hfma2(*reinterpret_cast(&hi), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); +} + +template <> +__device__ inline void dequant(int q, + half2* frag_b) { + dequant(q, frag_b); +} + +template <> +__device__ inline void dequant(int q, + half2* frag_b) { + const int LO = 0x000f000f; + const int HI = 0x00f000f0; + const int EX = 0x64006400; + // Guarantee that the `(a & b) | c` operations are LOP3s. + // clang-format off + int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); + int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); + // clang-format on + // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point + // directly into `SUB` and `ADD`. + const int SUB = 0x64006400; + const int MUL = 0x2c002c00; + const int ADD = 0xd400d400; + frag_b[0] = __hsub2(*reinterpret_cast(&lo), + *reinterpret_cast(&SUB)); + frag_b[1] = __hfma2(*reinterpret_cast(&hi), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); +} + +template <> +__device__ inline void dequant( + int q, nv_bfloat162* frag_b) { + static constexpr uint32_t MASK = 0x000f000f; + static constexpr uint32_t EX = 0x43004300; + + // Guarantee that the `(a & b) | c` operations are LOP3s. + // clang-format off + int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); + q >>= 4; + int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); + // clang-format on + + frag_b[0] = *reinterpret_cast(&lo); + frag_b[1] = *reinterpret_cast(&hi); +} + +template <> +__device__ inline void dequant( + int q, nv_bfloat162* frag_b) { + dequant(q, frag_b); + + static constexpr uint32_t SUB = 0x43084308; + + frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&SUB)); + frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&SUB)); +} + +template <> +__device__ inline void dequant( + int q, nv_bfloat162* frag_b) { + dequant(q, frag_b); +} + +template <> +__device__ inline void dequant( + int q, nv_bfloat162* frag_b) { + dequant(q, frag_b); + + static constexpr uint32_t SUB = 0x43004300; + + frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&SUB)); + frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&SUB)); +} + +// +// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or +// bf16 Reference: +// - FP16: +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85 +// - BF16: +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175 +// +template <> +__device__ inline void dequant(int q, + half2* frag_b) { + static constexpr uint32_t mask_for_elt_01 = 0x5250; + static constexpr uint32_t mask_for_elt_23 = 0x5351; + static constexpr uint32_t start_byte_for_fp16 = 0x64646464; + + uint32_t lo = prmt(q); + uint32_t hi = prmt(q); + + frag_b[0] = *reinterpret_cast(&lo); + frag_b[1] = *reinterpret_cast(&hi); +} + +template <> +__device__ inline void dequant( + int q, half2* frag_b) { + dequant(q, frag_b); + + static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480; + frag_b[0] = __hsub2(frag_b[0], + *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); + frag_b[1] = __hsub2(frag_b[1], + *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); +} + +template <> +__device__ inline void dequant(int q, + half2* frag_b) { + dequant(q, frag_b); +} + +template <> +__device__ inline void dequant(int q, + half2* frag_b) { + dequant(q, frag_b); + + static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400; + frag_b[0] = __hsub2(frag_b[0], + *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); + frag_b[1] = __hsub2(frag_b[1], + *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); +} + +template <> +__device__ inline void dequant( + int q, nv_bfloat162* frag_b) { + float fp32_intermediates[4]; + uint32_t* fp32_intermediates_casted = + reinterpret_cast(fp32_intermediates); + + static constexpr uint32_t fp32_base = 0x4B000000; + fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); + fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); + fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); + fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); + + fp32_intermediates[0] -= 8388736.f; + fp32_intermediates[1] -= 8388736.f; + fp32_intermediates[2] -= 8388736.f; + fp32_intermediates[3] -= 8388736.f; + + uint32_t* bf16_result_ptr = reinterpret_cast(frag_b); + bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], + fp32_intermediates_casted[1], 0x7632); + bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], + fp32_intermediates_casted[3], 0x7632); +} + +template <> +__device__ inline void dequant( + int q, nv_bfloat162* frag_b) { + float fp32_intermediates[4]; + uint32_t* fp32_intermediates_casted = + reinterpret_cast(fp32_intermediates); + + static constexpr uint32_t fp32_base = 0x4B000000; + fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); + fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); + fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); + fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); + + fp32_intermediates[0] -= 8388608.f; + fp32_intermediates[1] -= 8388608.f; + fp32_intermediates[2] -= 8388608.f; + fp32_intermediates[3] -= 8388608.f; + + uint32_t* bf16_result_ptr = reinterpret_cast(frag_b); + bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], + fp32_intermediates_casted[1], 0x7632); + bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], + fp32_intermediates_casted[3], 0x7632); +} + +template <> +__device__ inline void dequant( + int q, half2* frag_b) { + // Constants for FP8 (E4M3) and FP16 formats + constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5; + constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT; + constexpr int MASK = 0x7F007F00; + + // Extract and shift FP8 values to FP16 format + int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); + q <<= 8; + int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); + + // Note: reverse indexing is intentional because weights are permuted + frag_b[1] = *reinterpret_cast(&Out1); + frag_b[0] = *reinterpret_cast(&Out2); +} + +template <> +__device__ inline void dequant( + int q, half2* frag_b) { + dequant(q, frag_b); + + // Constants for FP8 (E4M3) and FP16 formats + constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5; + + // Construct and apply exponent bias + constexpr int BIAS_OFFSET = + (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); + const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); + + // Convert to half2 and apply bias + frag_b[1] = __hmul2(frag_b[1], bias_reg); + frag_b[0] = __hmul2(frag_b[0], bias_reg); +} + +template <> +__device__ inline void dequant( + int q, nv_bfloat162* frag_b) { + // Constants for FP8 (E4M3) and BF16 formats + constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; + constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; + + constexpr int MASK = 0x7F007F00; + + // Extract and shift FP8 values to BF16 format + int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); + q <<= 8; + int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); + + // Note: reverse indexing is intentional because weights are permuted + frag_b[1] = *reinterpret_cast(&Out1); + frag_b[0] = *reinterpret_cast(&Out2); +} + +template <> +__device__ inline void dequant( + int q, nv_bfloat162* frag_b) { + dequant(q, frag_b); + + // Constants for FP8 (E4M3) and BF16 formats + constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; + + // Construct and apply exponent bias + constexpr int BIAS_OFFSET = + (1 << (BF16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); + // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent + // position + constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; + const nv_bfloat162 bias_reg = + __float2bfloat162_rn(*reinterpret_cast(&BIAS)); + + // Convert to bfloat162 and apply bias + frag_b[1] = __hmul2(frag_b[1], bias_reg); + frag_b[0] = __hmul2(frag_b[0], bias_reg); +} + +template <> +__device__ inline void dequant(int q, + half2* frag_b) { + // Constants for FP4 (E2M1) and FP16 formats + constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5; + constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP4_EXPONENT; + constexpr int MASK = 0x70007000; + + // Extract and shift FP4 values to FP16 format + int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); + q <<= 4; + int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); + + // Note: reverse indexing is intentional because weights are permuted + frag_b[1] = *reinterpret_cast(&Out1); + frag_b[0] = *reinterpret_cast(&Out2); +} + +template <> +__device__ inline void dequant( + int q, half2* frag_b) { + dequant(q, frag_b); + + // Constants for FP4 (E2M1) and FP16 formats + constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5; + + // Construct and apply exponent bias + constexpr int BIAS_OFFSET = + (1 << (FP16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1)); + const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); + + // Convert to half2 and apply bias + frag_b[1] = __hmul2(frag_b[1], bias_reg); + frag_b[0] = __hmul2(frag_b[0], bias_reg); +} + +template <> +__device__ inline void dequant( + int q, nv_bfloat162* frag_b) { + // Constants for FP4 (E2M1) and FP16 formats + constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8; + constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP4_EXPONENT; + constexpr int MASK = 0x70007000; + + // Extract and shift FP4 values to FP16 format + int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); + q <<= 4; + int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); + + // Note: reverse indexing is intentional because weights are permuted + frag_b[1] = *reinterpret_cast(&Out1); + frag_b[0] = *reinterpret_cast(&Out2); +} + +template <> +__device__ inline void dequant( + int q, nv_bfloat162* frag_b) { + dequant(q, frag_b); + + // Constants for FP4 (E2M1) and BF16 formats + constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8; + + // Construct and apply exponent bias + constexpr int BIAS_OFFSET = + (1 << (BF16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1)); + // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent + // position + constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; + const nv_bfloat162 bias_reg = + __float2bfloat162_rn(*reinterpret_cast(&BIAS)); + + // Convert to half2 and apply bias + frag_b[1] = __hmul2(frag_b[1], bias_reg); + frag_b[0] = __hmul2(frag_b[0], bias_reg); +} + +template <> +__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kFE2M1f.id(), true>( + int q, __nv_fp8x4_e4m3* frag_b) { + // Constants for FP4 (E2M1) and FP16 formats + constexpr int FP4_EXPONENT = 2, FP8_EXPONENT = 4; + constexpr int RIGHT_SHIFT = FP8_EXPONENT - FP4_EXPONENT; + constexpr int MASK = 0x70707070; + + // Extract and shift FP4 values to FP16 format + int Out1 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT); + q <<= 4; + int Out2 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT); + + // Note1: reverse indexing is intentional because weights are permuted + // Note2: when dequant to 8bit type, we write to `frag_b[2]` instead of + // `frag_b[1]` to fit the layout of tensorcore + frag_b[1] = *reinterpret_cast(&Out1); + frag_b[0] = *reinterpret_cast(&Out2); +} + +template <> +__device__ inline void dequant( + int q, int32_t* frag_b) { + constexpr int repeated_zp = 0x08080808; + constexpr int MASK = 0x80808080; + + frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK; + q >>= 4; + frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK; +} + +template <> +__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kU4B8.id(), true>( + int q, __nv_fp8x4_e4m3* frag_b) { + int s = q & 0x08080808; + int Out1 = ((q & 0x07070707) | (s << 4)) + (s >> 3); + q >>= 4; + s = q & 0x08080808; + int Out2 = ((q & 0x07070707) | (s << 4)) + (s >> 3); + + frag_b[0] = *reinterpret_cast(&Out1); + frag_b[1] = *reinterpret_cast(&Out2); +} + +template +__device__ inline void dequant_fp8_scales(int q, scalar_t2* frag_b); + +template <> +__device__ inline void dequant_fp8_scales( + int q, half2* frag_b) { + int Out1 = (q & 0xFF00FF00) >> 1; + ; + q <<= 8; + int Out2 = (q & 0xFF00FF00) >> 1; + + // Note: reverse indexing is intentional because weights are permuted + frag_b[1] = *reinterpret_cast(&Out1); + frag_b[0] = *reinterpret_cast(&Out2); +}; + +template <> +__device__ inline void dequant_fp8_scales( + int q, nv_bfloat162* frag_b) { + constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; + constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; + constexpr int MASK = 0x7F007F00; + + // Extract and shift FP8 values to BF16 format + int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); + q <<= 8; + int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); + + // Note: reverse indexing is intentional because weights are permuted + frag_b[1] = *reinterpret_cast(&Out1); + frag_b[0] = *reinterpret_cast(&Out2); +} + +template <> +__device__ inline void dequant_fp8_scales( + int q, nv_bfloat162* frag_b) { + // In this conversion, 2 ** -127 in FP8E8M0 would become 0 in BF16, + // but we assume that such a extreme value would not occur in real models. + int Out1 = (q & 0xFF00FF00) >> 1; + q <<= 7; + int Out2 = q & 0x7F807F80; + + // Note: reverse indexing is intentional because weights are permuted + frag_b[1] = *reinterpret_cast(&Out1); + frag_b[0] = *reinterpret_cast(&Out2); +}; + +// subtract zero point in quanted format and then dequant +template +__device__ inline void sub_zp_and_dequant(int q, scalar_t2* frag_b, int zp); + +template <> +__device__ inline void sub_zp_and_dequant( + int q, int32_t* frag_b, int zp) { + // INT4 with zp -> INT8 + // see https://github.com/vllm-project/vllm/pull/24722 + int repeated_zp = 0x01010101 * zp; + int MASK = 0x80808080; + + frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK; + q >>= 4; + frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK; +} + +template <> +__device__ inline void sub_zp_and_dequant<__nv_fp8x4_e4m3, vllm::kU4.id(), + true>(int q, __nv_fp8x4_e4m3* frag_b, + int zp) { + // INT4 with zp -> FP8 + // see https://github.com/vllm-project/vllm/pull/24722 + uint32_t u_q = *reinterpret_cast(&q); + uint32_t u_zp = *reinterpret_cast(&zp); + uint32_t u_zp1 = u_zp + 1; + uint32_t repeated_zp = 0x01010101 * u_zp; + + uint32_t q0, s; + q0 = (u_q & 0x0F0F0F0F) | 0x70707070; + s = (q0 + repeated_zp) & 0x80808080; + uint32_t Out1 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s; + + u_q >>= 4; + q0 = (u_q & 0x0F0F0F0F) | 0x70707070; + s = (q0 + repeated_zp) & 0x80808080; + uint32_t Out2 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s; + + frag_b[0] = *reinterpret_cast(&Out1); + frag_b[1] = *reinterpret_cast(&Out2); +} + +#endif + +} // namespace MARLIN_NAMESPACE_NAME diff --git a/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu new file mode 100644 index 000000000..d268421a8 --- /dev/null +++ b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu @@ -0,0 +1,373 @@ +// GPTQ-Marlin weight-repack CUDA op, vendored from vLLM (Apache-2.0). +// +// Source: vllm-project/vllm +// csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu +// The __global__ ``gptq_marlin_repack_kernel`` is copied VERBATIM; only the host +// wrapper + op registration are rewritten from vLLM's ``torch::stable`` ABI to +// the classic ``torch::Tensor`` + ``TORCH_LIBRARY`` ABI mstar's vendored ops use +// (see utils/fused_moe/csrc/moe_align_block_size.cu). W4A16 only: the ``is_a_8bit`` +// (W4A8) template arm is dropped (fixed to false), so no fp8/int8 activation path. +// Registered under ``_mstar_marlin_C`` so it never collides with a real vLLM ``_C``. +// +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include "marlin.cuh" + +#include +#include +#include +#include + +namespace marlin { + +template +__global__ void gptq_marlin_repack_kernel( + uint32_t const* __restrict__ b_q_weight_ptr, + uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr, + int size_k, int size_n) { + constexpr int pack_factor = 32 / num_bits; + + constexpr int target_tile_n_size = tile_n_size / (is_a_8bit ? 2 : 1); + constexpr int target_tile_k_size = tile_k_size * (is_a_8bit ? 2 : 1); + int k_tiles = size_k / target_tile_k_size; + int n_tiles = size_n / target_tile_n_size; + int block_k_tiles = div_ceil(k_tiles, gridDim.x); + + auto start_k_tile = blockIdx.x * block_k_tiles; + if (start_k_tile >= k_tiles) { + return; + } + + int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); + + // Wait until the next thread tile has been loaded to shared memory. + auto wait_for_stage = [&]() { + // We only have `stages - 2` active fetches since we are double buffering + // and can only issue the next fetch when it is guaranteed that the previous + // shared memory load is fully complete (as it may otherwise be + // overwritten). + cp_async_wait(); + __syncthreads(); + }; + + extern __shared__ int4 sh[]; + + constexpr int perm_size = target_tile_k_size / 4; + + int4* sh_perm_ptr = sh; + int4* sh_pipe_ptr = sh_perm_ptr; + if constexpr (has_perm) { + sh_pipe_ptr += perm_size; + } + + constexpr int tile_ints = target_tile_k_size / pack_factor; + + constexpr int stage_n_threads = target_tile_n_size / 4; + constexpr int stage_k_threads = has_perm ? target_tile_k_size : tile_ints; + constexpr int stage_size = stage_k_threads * stage_n_threads; + + auto load_perm_to_shared = [&](int k_tile_id) { + int first_k_int4 = (k_tile_id * target_tile_k_size) / 4; + + int4 const* perm_int4_ptr = reinterpret_cast(perm_ptr); + + if (threadIdx.x < perm_size) { + sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x]; + } + __syncthreads(); + }; + + auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { + if (n_tile_id >= n_tiles) { + cp_async_fence(); + return; + } + + int first_n = n_tile_id * target_tile_n_size; + + int4* sh_ptr = sh_pipe_ptr + stage_size * pipe; + + if constexpr (has_perm) { + if (threadIdx.x < stage_size) { + auto k_id = threadIdx.x / stage_n_threads; + auto n_id = threadIdx.x % stage_n_threads; + + uint32_t const* sh_perm_int_ptr = + reinterpret_cast(sh_perm_ptr); + + int src_k = sh_perm_int_ptr[k_id]; + int src_k_packed = src_k / pack_factor; + + cp_async4( + &sh_ptr[k_id * stage_n_threads + n_id], + reinterpret_cast(&( + b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)]))); + } + + } else { + if (threadIdx.x < stage_size) { + auto k_id = threadIdx.x / stage_n_threads; + auto n_id = threadIdx.x % stage_n_threads; + + int first_k = k_tile_id * target_tile_k_size; + int first_k_packed = first_k / pack_factor; + + cp_async4(&sh_ptr[k_id * stage_n_threads + n_id], + reinterpret_cast( + &(b_q_weight_ptr[(first_k_packed + k_id) * size_n + + first_n + (n_id * 4)]))); + } + } + + cp_async_fence(); + }; + + auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { + if (n_tile_id >= n_tiles) { + return; + } + + auto warp_id = threadIdx.x / 32; + auto th_id = threadIdx.x % 32; + + if (warp_id >= 4) { + return; + } + + int tc_col = th_id / 4; + int tc_row = (th_id % 4) * (is_a_8bit ? 4 : 2); + + constexpr int tc_offsets[4] = {0, 1, 8, 9}; + + int cur_n = (warp_id / (is_a_8bit ? 2 : 1)) * 16 + tc_col; + + constexpr int sh_stride = target_tile_n_size; + constexpr uint32_t mask = (1 << num_bits) - 1; + + int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe; + uint32_t* sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); + + uint32_t* sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); + + uint32_t vals[8]; + + if constexpr (has_perm) { + static_assert(!is_a_8bit); + for (int i = 0; i < 4; i++) { + int k_idx = tc_row + tc_offsets[i]; + + uint32_t src_k = sh_perm_int_ptr[k_idx]; + uint32_t src_k_pos = src_k % pack_factor; + + uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n]; + uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask; + + uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8]; + uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask; + + vals[i] = b1_cur_val; + vals[4 + i] = b2_cur_val; + } + + } else { + uint32_t b1_vals[tile_ints]; + uint32_t b2_vals[tile_ints]; + +#pragma unroll + for (int i = 0; i < tile_ints; i++) { + if constexpr (is_a_8bit) { + b1_vals[i] = + sh_stage_int_ptr[cur_n + sh_stride * i + (warp_id % 2) * 8]; + } else { + b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i]; + b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i]; + } + } + +#pragma unroll + for (int i = 0; i < 4; i++) { + int cur_elem = tc_row + (is_a_8bit ? i : tc_offsets[i]); + int cur_int = cur_elem / pack_factor; + int cur_pos = cur_elem % pack_factor; + + vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask; + if constexpr (is_a_8bit) + vals[4 + i] = + (b1_vals[cur_int + tile_ints / 2] >> (cur_pos * num_bits)) & mask; + else + vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask; + } + } + + constexpr int tile_size = + target_tile_k_size * target_tile_n_size / pack_factor; + int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size; + + // Result of: + // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h + if constexpr (!is_a_8bit && num_bits == 4) { + int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; + + uint32_t res = 0; +#pragma unroll + for (int i = 0; i < 8; i++) { + res |= vals[pack_idx[i]] << (i * 4); + } + + out_ptr[out_offset + th_id * 4 + warp_id] = res; + + } else if constexpr (is_a_8bit && num_bits == 4) { + int pack_idx[8] = {0, 4, 1, 5, 2, 6, 3, 7}; + + uint32_t res = 0; +#pragma unroll + for (int i = 0; i < 8; i++) { + res |= vals[pack_idx[i]] << (i * 4); + } + + out_ptr[out_offset + th_id * 4 + warp_id] = res; + + } else { + constexpr int pack_idx[4] = {0, 2, 1, 3}; + + uint32_t res1 = 0; + uint32_t res2 = 0; +#pragma unroll + for (int i = 0; i < 4; i++) { + const int ii = is_a_8bit ? i : pack_idx[i]; + res1 |= vals[ii] << (i * 8); + res2 |= vals[4 + ii] << (i * 8); + } + + out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; + out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; + } + }; + + auto start_pipes = [&](int k_tile_id, int n_tile_id) { +#pragma unroll + for (int pipe = 0; pipe < repack_stages - 1; pipe++) { + fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); + } + + wait_for_stage(); + }; +#pragma unroll + for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) { + int n_tile_id = 0; + + if constexpr (has_perm) { + load_perm_to_shared(k_tile_id); + } + + start_pipes(k_tile_id, n_tile_id); + + while (n_tile_id < n_tiles) { +#pragma unroll + for (int pipe = 0; pipe < repack_stages; pipe++) { + fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, + n_tile_id + pipe + repack_stages - 1); + repack_tile(pipe, k_tile_id, n_tile_id + pipe); + wait_for_stage(); + } + n_tile_id += repack_stages; + } + } +} + +} // namespace marlin + +// W4A16 only: is_a_8bit is fixed to false, so only the has_perm ∈ {false, true} +// arms are instantiated (act-order not exercised by Kimi, but kept for reuse). +#define CALL_IF(NUM_BITS, HAS_PERM) \ + else if (num_bits == NUM_BITS && has_perm == HAS_PERM) { \ + cudaFuncSetAttribute( \ + marlin::gptq_marlin_repack_kernel, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ + marlin::gptq_marlin_repack_kernel \ + <<>>( \ + b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \ + } + +torch::Tensor gptq_marlin_repack(torch::Tensor b_q_weight, torch::Tensor perm, + int64_t size_k, int64_t size_n, + int64_t num_bits) { + // Verify compatibility with marlin tile of 16x64 + TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k, + " is not divisible by tile_k_size = ", marlin::tile_k_size); + TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n, + " is not divisible by tile_n_size = ", marlin::tile_n_size); + + TORCH_CHECK(num_bits == 4 || num_bits == 8, + "num_bits must be 4 or 8. Got = ", num_bits); + int const pack_factor = 32 / num_bits; + + // Verify B + TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0), + "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0), + ", size_k = ", size_k, ", pack_factor = ", pack_factor); + TORCH_CHECK(b_q_weight.size(1) == size_n, + "b_q_weight.size(1) = ", b_q_weight.size(1), + " is not size_n = ", size_n); + + // Verify device and strides + TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU"); + TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); + TORCH_CHECK(b_q_weight.scalar_type() == torch::kInt, + "b_q_weight type is not kInt"); + + TORCH_CHECK(perm.is_cuda(), "perm is not on GPU"); + TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); + TORCH_CHECK(perm.scalar_type() == torch::kInt, "perm type is not kInt"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const int device_index = b_q_weight.get_device(); + + // Alloc buffers + torch::Tensor out = torch::empty( + {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor}, + b_q_weight.options()); + + // Detect if there is act_order + bool has_perm = perm.size(0) != 0; + + // Get ptrs + uint32_t const* b_q_weight_ptr = + reinterpret_cast(b_q_weight.data_ptr()); + uint32_t const* perm_ptr = + reinterpret_cast(perm.data_ptr()); + uint32_t* out_ptr = reinterpret_cast(out.data_ptr()); + + int blocks; + cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index); + + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index); + TORCH_CHECK(max_shared_mem > 0); + + if (false) { + } + CALL_IF(4, false) + CALL_IF(4, true) + CALL_IF(8, false) + CALL_IF(8, true) + else { + TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, + ", has_perm = ", has_perm); + } + + return out; +} + +// The op schemas for the whole ``_mstar_marlin_C`` namespace are declared by the +// single TORCH_LIBRARY block in marlin_moe.cu (one per namespace per extension); +// this file only provides the repack impl. +TORCH_LIBRARY_IMPL(_mstar_marlin_C, CUDA, m) { + m.impl("gptq_marlin_repack", &gptq_marlin_repack); +} diff --git a/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin.cuh b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin.cuh new file mode 100644 index 000000000..da8f39f53 --- /dev/null +++ b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin.cuh @@ -0,0 +1,182 @@ +// Marlin common device header, vendored VERBATIM from vLLM (Apache-2.0). +// +// Source: vllm-project/vllm csrc/libtorch_stable/quantization/marlin/marlin.cuh +// Self-contained: includes only + . Provides the Marlin +// tile/thread constants and the cp.async helpers used by the repack and GEMM +// kernels. Kept byte-identical to upstream so future syncs are a clean diff. +// +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +#ifndef _marlin_cuh + #define _marlin_cuh + #include + #include + #include + #include + + #ifndef MARLIN_NAMESPACE_NAME + #define MARLIN_NAMESPACE_NAME marlin + #endif + +namespace MARLIN_NAMESPACE_NAME { + +// Marlin params + +// 8 warps are a good choice since every SM has 4 schedulers and having more +// than 1 warp per schedule allows some more latency hiding. At the same time, +// we want relatively few warps to have many registers per warp and small tiles. +static constexpr int default_threads = 256; + +static constexpr int pipe_stages = + 4; // 4 pipeline stages fit into shared memory + +static constexpr int min_thread_n = 64; +static constexpr int min_thread_k = 64; +static constexpr int max_thread_n = 256; + +static constexpr int tile_size = 16; +static constexpr int max_par = 16; + +// Repack params +static constexpr int repack_stages = 8; + +static constexpr int repack_threads = 256; + +static constexpr int tile_k_size = tile_size; +static constexpr int tile_n_size = tile_k_size * 4; + +// Helpers +template +struct Vec { + T elems[n]; + __device__ T& operator[](int i) { return elems[i]; } +}; + +using I4 = Vec; + +constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; } + + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + +__device__ inline void cp_async1_ca_pred(void* smem_ptr, const void* glob_ptr, + bool pred = true) { + if (pred) { + reinterpret_cast(smem_ptr)[0] = + reinterpret_cast(glob_ptr)[0]; + } +} + +__device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr, + bool pred = true) { + if (pred) { + reinterpret_cast(smem_ptr)[0] = + reinterpret_cast(glob_ptr)[0]; + } +} + +__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr, + bool pred = true) { + if (pred) { + reinterpret_cast(smem_ptr)[0] = + reinterpret_cast(glob_ptr)[0]; + } +} + +__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr, + bool pred = true) { + if (pred) { + reinterpret_cast(smem_ptr)[0] = + reinterpret_cast(glob_ptr)[0]; + } +} + +__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) { + reinterpret_cast(smem_ptr)[0] = + reinterpret_cast(glob_ptr)[0]; +} + +__device__ inline void cp_async_fence() {} + +template +__device__ inline void cp_async_wait() {} + + #else + +__device__ inline void cp_async1_ca_pred(void* smem_ptr, const void* glob_ptr, + bool pred = true) { + const int BYTES = 4; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.ca.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int)pred), + "r"(smem), "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr, + bool pred = true) { + const int BYTES = 8; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.ca.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int)pred), + "r"(smem), "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr, + bool pred = true) { + const int BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.ca.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int)pred), + "r"(smem), "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr, + bool pred = true) { + const int BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int)pred), + "r"(smem), "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) { + const int BYTES = 16; + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "{\n" + " cp.async.cg.shared.global [%0], [%1], %2;\n" + "}\n" ::"r"(smem), + "l"(glob_ptr), "n"(BYTES)); +} + +__device__ inline void cp_async_fence() { + asm volatile("cp.async.commit_group;\n" ::); +} + +template +__device__ inline void cp_async_wait() { + asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); +} + + #endif + +} // namespace MARLIN_NAMESPACE_NAME + +#endif diff --git a/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh new file mode 100644 index 000000000..a4807a688 --- /dev/null +++ b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh @@ -0,0 +1,149 @@ + +#ifndef _data_types_cuh +#define _data_types_cuh +#include "marlin.cuh" +#include "core/scalar_type.hpp" +#include +#include +#include + +#ifndef MARLIN_NAMESPACE_NAME + #define MARLIN_NAMESPACE_NAME marlin +#endif + +namespace MARLIN_NAMESPACE_NAME { + +template +class MarlinScalarType {}; + +template <> +class MarlinScalarType { + public: + using scalar_t = half; + using scalar_t2 = half2; + using scalar_t4 = half2; + using scalar_32bit_t = half2; + + // Matrix fragments for tensor core instructions; their precise layout is + // documented here: + // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type + using FragA = Vec; + using FragB = Vec; + using FragC = Vec; + using FragS = Vec; + using FragS0 = Vec<__nv_fp8x2_e4m3, 1>; + using FragZP = Vec; + + static __device__ float inline num2float(const half x) { + return __half2float(x); + } + + static __device__ half2 inline num2num2(const half x) { + return __half2half2(x); + } + + static __device__ half2 inline nums2num2(const half x1, const half x2) { + return __halves2half2(x1, x2); + } + + static __host__ __device__ half inline float2num(const float x) { + return __float2half(x); + } + + static __host__ __device__ float2 inline num22float2(const half2 x) { + return __half22float2(x); + } +}; + +template <> +class MarlinScalarType { + public: + using scalar_t = nv_bfloat16; + using scalar_t2 = nv_bfloat162; + using scalar_t4 = nv_bfloat162; + using scalar_32bit_t = nv_bfloat162; + + using FragA = Vec; + using FragB = Vec; + using FragC = Vec; + using FragS = Vec; + using FragS0 = Vec<__nv_fp8x2_e4m3, 1>; + using FragZP = Vec; + +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 + static __device__ float inline num2float(const nv_bfloat16 x) { + return __bfloat162float(x); + } + + static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) { + return __bfloat162bfloat162(x); + } + + static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1, + const nv_bfloat16 x2) { + return __halves2bfloat162(x1, x2); + } + + static __host__ __device__ nv_bfloat16 inline float2num(const float x) { + return __float2bfloat16(x); + } + + static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) { + return __bfloat1622float2(x); + } +#endif +}; + +template <> +class MarlinScalarType { + public: + using scalar_t = __nv_fp8_e4m3; + using scalar_t2 = __nv_fp8x2_e4m3; + using scalar_t4 = __nv_fp8x4_e4m3; + using scalar_32bit_t = __nv_fp8x4_e4m3; + + using FragA = Vec<__nv_fp8x4_e4m3, 4>; + using FragB = Vec<__nv_fp8x4_e4m3, 2>; + using FragC = Vec; + using FragZP = Vec<__nv_fp8x2_e4m3, 4>; + + static __host__ __device__ + float2 inline num22float2(const __nv_fp8x2_e4m3 x) { + return (float2)x; + } +}; + +template <> +class MarlinScalarType { + public: + using scalar_t = int8_t; + using scalar_t2 = int16_t; + using scalar_t4 = int32_t; + using scalar_32bit_t = int32_t; + + using FragA = Vec; + using FragB = Vec; + using FragC = Vec; + using FragZP = Vec; +}; + +template +class MarlinScalarType2 {}; + +template <> +class MarlinScalarType2 : public MarlinScalarType {}; + +template <> +class MarlinScalarType2 + : public MarlinScalarType {}; + +template <> +class MarlinScalarType2<__nv_fp8_e4m3> + : public MarlinScalarType {}; + +template <> +class MarlinScalarType2 : public MarlinScalarType {}; + +} // namespace MARLIN_NAMESPACE_NAME + +#endif diff --git a/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_mma.h b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_mma.h new file mode 100644 index 000000000..6ec2aaafc --- /dev/null +++ b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_mma.h @@ -0,0 +1,269 @@ + +#include "marlin_dtypes.cuh" + +namespace MARLIN_NAMESPACE_NAME { + +// m16n8k16 tensor core mma instruction with fp16 inputs and fp32 +// output/accumulation. +template +__device__ inline void mma( + const typename MarlinScalarType::FragA& a_frag, + const typename MarlinScalarType::FragB& frag_b, + typename MarlinScalarType::FragC& frag_c, int idx = 0) { + const uint32_t* a = reinterpret_cast(&a_frag); + const uint32_t* b = reinterpret_cast(&frag_b); + using scalar_t = typename MarlinScalarType::scalar_t; + if constexpr (!std::is_same::value || k_size != 16) { + static_assert(!use_fp16_accum); + } + + if constexpr (k_size == 16) { + if constexpr (std::is_same::value && !use_fp16_accum) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 + float* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(b[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]), + "f"(c[3])); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[2]), "r"(a[3]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), + "f"(c[3])); +#else + float* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); +#endif + } else if constexpr (std::is_same::value && + use_fp16_accum) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 + uint32_t* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " + "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n" + : "=r"(c[0]), "=r"(c[1]) + : "r"(a[0]), "r"(a[1]), "r"(b[0]), "r"(c[0]), "r"(c[1])); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " + "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n" + : "=r"(c[0]), "=r"(c[1]) + : "r"(a[2]), "r"(a[3]), "r"(b[1]), "r"(c[0]), "r"(c[1])); +#else + uint32_t* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n" + : "=r"(c[0]), "=r"(c[1]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "r"(c[0]), "r"(c[1])); +#endif + } else if constexpr (std::is_same::value) { + float* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } else if constexpr (std::is_same::value) { + float* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 " + "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "f"(c[0]), + "f"(c[1]), "f"(c[2]), "f"(c[3])); + } else if constexpr (std::is_same::value) { + int32_t* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite " + "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" + : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) + : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "r"(c[0]), + "r"(c[1]), "r"(c[2]), "r"(c[3])); + } + } else if (k_size == 32) { + if constexpr (std::is_same::value) { + float* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } else if constexpr (std::is_same::value) { + int32_t* c = reinterpret_cast(&frag_c); +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " + "{%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(c[0]), "=r"(c[1]) + : "r"(a[0]), "r"(b[0]), "r"(c[0]), "r"(c[1])); + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " + "{%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(c[2]), "=r"(c[3]) + : "r"(a[1]), "r"(b[0]), "r"(c[2]), "r"(c[3])); + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " + "{%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(c[0]), "=r"(c[1]) + : "r"(a[2]), "r"(b[1]), "r"(c[0]), "r"(c[1])); + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " + "{%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(c[2]), "=r"(c[3]) + : "r"(a[3]), "r"(b[1]), "r"(c[2]), "r"(c[3])); +#else + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3])); +#endif + } + } +} + +template +__device__ inline void mma_trans( + const typename MarlinScalarType::FragA& a_frag, + const typename MarlinScalarType::FragB& frag_b, + const typename MarlinScalarType::FragB& frag_b2, + typename MarlinScalarType::FragC& frag_c) { + const uint32_t* a = reinterpret_cast(&a_frag); + const uint32_t* b = reinterpret_cast(&frag_b); + const uint32_t* b2 = reinterpret_cast(&frag_b2); + float* c = reinterpret_cast(&frag_c); + using scalar_t = typename MarlinScalarType::scalar_t; + if constexpr (!std::is_same::value || k_size != 16) { + static_assert(!use_fp16_accum); + } + + if constexpr (k_size == 16) { + if constexpr (std::is_same::value && !use_fp16_accum) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 + float* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]), + "f"(c[3])); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(b[1]), "r"(b2[1]), "r"(a[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), + "f"(c[3])); +#else + float* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); +#endif + } else if constexpr (std::is_same::value && + use_fp16_accum) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 + uint32_t* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " + "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n" + : "=r"(c[0]), "=r"(c[1]) + : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1])); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " + "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n" + : "=r"(c[0]), "=r"(c[1]) + : "r"(b[1]), "r"(b2[1]), "r"(a[1]), "r"(c[0]), "r"(c[1])); +#else + uint32_t* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n" + : "=r"(c[0]), "=r"(c[1]) + : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), + "r"(c[0]), "r"(c[1])); +#endif + } else if constexpr (std::is_same::value) { + float* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } else if constexpr (std::is_same::value) { + float* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 " + "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]), + "f"(c[3])); + } else if constexpr (std::is_same::value) { + int32_t* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite " + "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" + : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) + : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]), "r"(c[2]), + "r"(c[3])); + } + } else { + if constexpr (std::is_same::value) { + float* c = reinterpret_cast(&frag_c); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) + : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); + } else if constexpr (std::is_same::value) { + int32_t* c = reinterpret_cast(&frag_c); +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " + "{%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(c[0]), "=r"(c[1]) + : "r"(b[0]), "r"(a[0]), "r"(c[0]), "r"(c[1])); + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " + "{%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(c[2]), "=r"(c[3]) + : "r"(b2[1]), "r"(a[0]), "r"(c[2]), "r"(c[3])); + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " + "{%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(c[0]), "=r"(c[1]) + : "r"(b[0]), "r"(a[1]), "r"(c[0]), "r"(c[1])); + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " + "{%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(c[2]), "=r"(c[3]) + : "r"(b2[1]), "r"(a[1]), "r"(c[2]), "r"(c[3])); +#else + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) + : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), + "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3])); +#endif + } + } +} + +} // namespace MARLIN_NAMESPACE_NAME \ No newline at end of file diff --git a/mstar/utils/marlin/loader.py b/mstar/utils/marlin/loader.py new file mode 100644 index 000000000..2d126b0bd --- /dev/null +++ b/mstar/utils/marlin/loader.py @@ -0,0 +1,81 @@ +"""JIT build + load of the vendored Marlin W4A16 CUDA ops. + +Mirrors :mod:`mstar.utils.fused_moe.align`: the Marlin CUDA sources (vendored +Apache-2.0 from vLLM under ``csrc/``) are JIT-compiled with +``torch.utils.cpp_extension.load`` on first use and registered as +``torch.ops._mstar_marlin_C.*`` — no ``vllm`` / ``sgl_kernel`` runtime dependency. + +If the build fails (no ``nvcc`` / no ``ninja`` / sm<80 / ABI mismatch) the loader +logs and returns ``False``; callers fall back to the Triton W4A16 path +(``fused_moe_kernel_w4a16``). Marlin is a *speed* layer over a correctness path +that already exists — exactly the CUDA-op-or-torch-fallback pattern ``align.py`` +uses for ``moe_align_block_size``. +""" +from __future__ import annotations + +import functools +import logging +import os + +import torch + +logger = logging.getLogger(__name__) + +_CSRC = os.path.join(os.path.dirname(__file__), "csrc") +_MARLIN = os.path.join(_CSRC, "libtorch_stable", "quantization", "marlin") +_MARLIN_MOE = os.path.join(_CSRC, "libtorch_stable", "moe", "marlin_moe_wna16") + +# Sources compiled into the ``_mstar_marlin_C`` extension: the repack op, the MoE +# GEMM host+device shim, and the pre-generated per-config kernel instantiations +# (``sm80_kernel_*.cu`` + ``kernel_selector.h``, produced once by the trimmed +# ``generate_kernels.py`` and vendored — GPTQ symmetric INT4, fp16/bf16 only). +_SOURCES = [ + os.path.join(_MARLIN, "gptq_marlin_repack.cu"), + os.path.join(_MARLIN_MOE, "marlin_moe.cu"), + os.path.join(_MARLIN_MOE, "sm80_kernel_bfloat16_u4b8_bfloat16.cu"), + os.path.join(_MARLIN_MOE, "sm80_kernel_float16_u4b8_float16.cu"), +] + +# Marlin's device code (cp.async, m16n8k16 MMA, bf16) requires sm80+. +_MIN_CAPABILITY = (8, 0) + + +@functools.lru_cache(maxsize=1) +def is_marlin_available() -> bool: + """JIT-build the Marlin ops once per process; return whether they are usable. + + Cached so compilation is attempted at most once. Any failure (missing + toolchain, unsupported GPU, compile/ABI error) is logged and the caller uses + the Triton W4A16 fallback. + """ + if not torch.cuda.is_available(): + return False + capability = torch.cuda.get_device_capability() + if capability < _MIN_CAPABILITY: + logger.warning( + "Marlin W4A16 needs sm%d%d+ (device is sm%d%d); using the Triton " + "in-kernel-dequant fallback.", + _MIN_CAPABILITY[0], _MIN_CAPABILITY[1], capability[0], capability[1], + ) + return False + try: + from torch.utils.cpp_extension import load + + load( + name="_mstar_marlin_C", + sources=list(_SOURCES), + is_python_module=False, + extra_include_paths=[_CSRC], + extra_cuda_cflags=["-O3", "-std=c++17", "--expt-relaxed-constexpr"], + verbose=False, + ) + # Touch an op so a registration failure surfaces here, not at call time. + _ = torch.ops._mstar_marlin_C.moe_wna16_marlin_gemm + return True + except Exception as e: # pragma: no cover -- depends on the build toolchain + logger.warning( + "Marlin W4A16: could not build the CUDA ops (%s); using the Triton " + "in-kernel-dequant fallback.", + e, + ) + return False diff --git a/mstar/utils/marlin/ops.py b/mstar/utils/marlin/ops.py new file mode 100644 index 000000000..c01988384 --- /dev/null +++ b/mstar/utils/marlin/ops.py @@ -0,0 +1,181 @@ +"""Python launchers over the vendored ``torch.ops._mstar_marlin_C`` Marlin ops. + +Mirrors vLLM's ``_custom_ops`` + ``marlin_utils`` helpers (same transformation +sequence: checkpoint INT4 → Marlin-repacked → GEMM), so the port is auditable +against the reference. Callers must gate on +:func:`mstar.utils.marlin.is_marlin_available` first — these dereference the JIT +op namespace directly. + +The routed-expert GEMM (:func:`fused_marlin_moe`) reuses mstar's existing MoE +plumbing — ``moe_align_block_size`` (token→expert sort), ``act_and_mul_triton`` +(SwiGLU), and ``moe_sum_reduce_triton`` (top-k fold) — so only the two INT4 +matmuls are Marlin; everything around them is shared with the bf16/Triton paths. +""" +from __future__ import annotations + +import torch + +from mstar.utils.fused_moe.align import moe_align_block_size +from mstar.utils.fused_moe.kernels import act_and_mul_triton, moe_sum_reduce_triton +from mstar.utils.marlin.scalar_type import UINT4B8_ID + +# Marlin repacked-weight tile (from the vendored marlin.cuh ``tile_size``). +_MARLIN_TILE = 16 + + +# --------------------------------------------------------------------------- +# Load-time repack (checkpoint GPTQ-packed INT4 → Marlin tiled layout) +# --------------------------------------------------------------------------- + +def gptq_marlin_repack( + b_q_weight: torch.Tensor, size_k: int, size_n: int, num_bits: int = 4 +) -> torch.Tensor: + """Repack a GPTQ-layout packed INT4 weight into Marlin tiled layout. + + ``b_q_weight`` is int32 ``(size_k // pack_factor, size_n)`` (K-major packed, + the layout Marlin's repack expects). Returns the Marlin-tiled int32 weight + ``(size_k // 16, size_n * 16 // pack_factor)``. + """ + perm = torch.empty(0, dtype=torch.int32, device=b_q_weight.device) + return torch.ops._mstar_marlin_C.gptq_marlin_repack( + b_q_weight, perm, size_k, size_n, num_bits + ) + + +def gptq_marlin_moe_repack( + b_q_weight: torch.Tensor, size_k: int, size_n: int, num_bits: int = 4 +) -> torch.Tensor: + """Per-expert :func:`gptq_marlin_repack` (mirrors vLLM's pure-Python loop). + + ``b_q_weight`` is int32 ``(E, size_k // pack_factor, size_n)``; returns + ``(E, size_k // 16, size_n * num_bits // 8)``. + """ + num_experts = b_q_weight.shape[0] + perm = torch.empty(0, dtype=torch.int32, device=b_q_weight.device) + output = torch.empty( + (num_experts, size_k // _MARLIN_TILE, size_n * (num_bits // 2)), + device=b_q_weight.device, + dtype=b_q_weight.dtype, + ) + for e in range(num_experts): + output[e] = torch.ops._mstar_marlin_C.gptq_marlin_repack( + b_q_weight[e], perm, size_k, size_n, num_bits + ) + return output + + +def _get_scale_perms() -> tuple[list[int], list[int]]: + """Marlin scale-permutation index tables (verbatim from vLLM marlin_utils).""" + scale_perm: list[int] = [] + for i in range(8): + scale_perm.extend([i + 8 * j for j in range(8)]) + scale_perm_single: list[int] = [] + for i in range(4): + scale_perm_single.extend([2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]]) + return scale_perm, scale_perm_single + + +def marlin_permute_scales( + s: torch.Tensor, size_k: int, size_n: int, group_size: int +) -> torch.Tensor: + """Permute a single expert's group scales into Marlin layout (vLLM parity).""" + scale_perm, scale_perm_single = _get_scale_perms() + if group_size < size_k and group_size != -1: + s = s.reshape((-1, len(scale_perm)))[:, scale_perm] + else: + s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single] + return s.reshape((-1, size_n)).contiguous() + + +def marlin_moe_permute_scales( + s: torch.Tensor, size_k: int, size_n: int, group_size: int +) -> torch.Tensor: + """Per-expert :func:`marlin_permute_scales`. ``s`` is ``(E, num_groups, size_n)``.""" + num_experts = s.shape[0] + output = torch.empty_like(s) + for e in range(num_experts): + output[e] = marlin_permute_scales(s[e], size_k, size_n, group_size) + return output + + +def marlin_make_workspace(device: torch.device, max_blocks_per_sm: int = 4) -> torch.Tensor: + """Marlin reduce/lock workspace: one int per (SM × max_blocks_per_sm).""" + sms = torch.cuda.get_device_properties(device).multi_processor_count + return torch.zeros(sms * max_blocks_per_sm, dtype=torch.int, device=device) + + +# --------------------------------------------------------------------------- +# Runtime: fused routed-expert Marlin GEMM +# --------------------------------------------------------------------------- + +def fused_marlin_moe( + hidden_states: torch.Tensor, + w1_marlin: torch.Tensor, + w2_marlin: torch.Tensor, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + workspace: torch.Tensor, + *, + activation: str = "silu", + reduce_results: bool = True, +) -> torch.Tensor: + """Marlin W4A16 routed-expert dispatch (gate_up GEMM → SwiGLU → down GEMM). + + Layout mirrors :func:`mstar.utils.fused_moe.fused_experts` but with Marlin + kernels: ``w1_marlin``/``w2_marlin`` are the Marlin-repacked int32 experts + (from :func:`gptq_marlin_moe_repack`), ``w1_scale``/``w2_scale`` the permuted + group scales (from :func:`marlin_moe_permute_scales`). Returns + ``(tokens, hidden)`` when ``reduce_results`` else the per-slot + ``(tokens, top_k, hidden)`` tensor the TP path all-reduces before folding. + """ + assert hidden_states.is_contiguous() and hidden_states.dim() == 2 + assert hidden_states.dtype in (torch.bfloat16, torch.float16) + M, K = hidden_states.shape + E = w1_marlin.shape[0] + top_k = topk_ids.shape[1] + # w2 marlin is (E, N//16, K*num_bits//8) -> intermediate size N = shape[1]*16. + N = w2_marlin.shape[1] * _MARLIN_TILE + + # Block-size selection (vLLM heuristic): smallest block that keeps ~>=0.9 + # expert occupancy. + block_size_m = 8 + for block_size_m in (8, 16, 32, 48, 64): + if M * top_k / E / block_size_m < 0.9: + break + + # Marlin requires fp32 combine weights and int32 expert ids. + topk_weights = topk_weights.to(torch.float32).contiguous() + topk_ids_i32 = topk_ids.to(torch.int32).contiguous() + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids_i32, block_size_m, E + ) + + # Gate+up GEMM: (M, K) x experts -> (M*top_k, 2N). + gate_up = torch.ops._mstar_marlin_C.moe_wna16_marlin_gemm( + hidden_states, None, w1_marlin, w1_scale, workspace, + sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights, + block_size_m, top_k, False, UINT4B8_ID, M, 2 * N, K, + True, False, True, + ) + + # SwiGLU: silu(gate) * up -> (M*top_k, N). + down_in = torch.empty((M * top_k, N), device=hidden_states.device, dtype=hidden_states.dtype) + act_and_mul_triton(gate_up, down_in, activation=activation) + + # Down GEMM (weighted): (M*top_k, N) x experts -> (M*top_k, K), routing + # weight folded in (mul_topk_weights=True). + down = torch.ops._mstar_marlin_C.moe_wna16_marlin_gemm( + down_in, None, w2_marlin, w2_scale, workspace, + sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights, + block_size_m, 1, True, UINT4B8_ID, M * top_k, K, N, + True, False, True, + ) + + cache3 = down.view(M, top_k, K) + if not reduce_results: + return cache3 + output = torch.empty_like(hidden_states) + moe_sum_reduce_triton(cache3, output, routed_scaling_factor=1.0) + return output diff --git a/mstar/utils/marlin/scalar_type.py b/mstar/utils/marlin/scalar_type.py new file mode 100644 index 000000000..b1a9c3e5f --- /dev/null +++ b/mstar/utils/marlin/scalar_type.py @@ -0,0 +1,18 @@ +"""vLLM ``ScalarType`` ids mirrored in Python for the Marlin ops. + +The vendored ``core/scalar_type.hpp`` packs a scalar type into a single int64 id +(field order exponent|mantissa|signed|bias|finite|nan_repr). The Marlin GEMM op +takes ``b_type_id`` and reconstructs the type via ``ScalarType::from_id``, so the +Python side must pass the exact same id. Only symmetric INT4 (``uint4b8``) is +built. The id is computed directly from the header's bit layout and re-validated +at runtime by the C++ op's ``TORCH_CHECK(b_type == kU4B8)``. + + uint4b8 = ScalarType::uint(size_bits=4, bias=8) + = (mantissa=4 << 8) | (bias=8 << 17) | (nan_repr=NAN_IEEE_754=1 << 50) +""" +from __future__ import annotations + +# vllm::kU4B8.id() — symmetric GPTQ-style INT4 (offset-binary, subtract 8). +UINT4B8_ID = (4 << 8) | (8 << 17) | (1 << 50) # == 1125899907892224 + +assert UINT4B8_ID == 1125899907892224, "uint4b8 id drifted from the vendored header" diff --git a/pyproject.toml b/pyproject.toml index 273a1c10f..0fb9d1b03 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -205,6 +205,12 @@ exclude = ["tests*", "benchmark*", "examples*"] # Ship the vendored CUDA source so it can be JIT-compiled on first use. [tool.setuptools.package-data] "mstar.utils.fused_moe" = ["csrc/*.cu", "csrc/*.cuh", "csrc/*.h", "csrc/*.cpp"] +# Marlin W4A16 sources live in a nested csrc tree (core/, libtorch_stable/...), +# including the pre-generated per-config kernel instantiations. +"mstar.utils.marlin" = [ + "csrc/**/*.cu", "csrc/**/*.cuh", "csrc/**/*.h", "csrc/**/*.hpp", + "csrc/**/*.cpp", "csrc/**/*.py", +] [tool.ruff] exclude = [".git", ".ruff_cache", ".venv", "ref"] diff --git a/test/integration/test_kimi_moe_marlin.py b/test/integration/test_kimi_moe_marlin.py new file mode 100644 index 000000000..f42540820 --- /dev/null +++ b/test/integration/test_kimi_moe_marlin.py @@ -0,0 +1,139 @@ +"""GPU golden for the Marlin backend wired into ``KimiSparseMoeBlock``. + +Where ``test_marlin_kernels.py`` exercises the kernels in isolation, this drives +the *block* wiring: ``process_weights_after_loading`` resolving the backend + +repacking the packed experts (and freeing the source packed params), the +``_use_marlin`` branch in ``forward`` passing fp32 combine weights, and the shared +expert / router riding alongside. The reference is the same block's bf16 path: +router + bf16 fused-expert GEMM on the dequantized experts + the (bf16) shared +expert. Only the routed-expert kernel differs, so a cosine floor + relative-L2 +bound is the correct gate (see ``test_marlin_kernels.py`` for why). + +Run: pytest test/integration/test_kimi_moe_marlin.py -v +""" +import pytest +import torch + +from mstar.model.kimi_k2_7._testing import fake_quantize_weight + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 0), + reason="Marlin-backed KimiSparseMoeBlock golden needs a CUDA GPU with sm80+", +) + +DEVICE = "cuda" +GROUP_SIZE = 32 +PACK_FACTOR = 8 + + +def _quantize_stack(weight): + E, N, K = weight.shape + packed = torch.empty((E, N, K // PACK_FACTOR), dtype=torch.int32, device=DEVICE) + scale = torch.empty((E, N, K // GROUP_SIZE), dtype=torch.bfloat16, device=DEVICE) + deq = torch.empty((E, N, K), dtype=torch.bfloat16, device=DEVICE) + for e in range(E): + p, s, d = fake_quantize_weight( + weight[e], num_bits=4, group_size=GROUP_SIZE, symmetric=True, + scale_dtype=torch.bfloat16, + ) + packed[e], scale[e], deq[e] = p.to(DEVICE), s.to(DEVICE), d.to(DEVICE) + return packed, scale, deq + + +def _build_block(): + """A tp=1 Marlin-legal KimiSparseMoeBlock, materialized on CUDA with random + router/shared weights and synthetic quantized routed experts loaded packed.""" + from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock + from mstar.model.kimi_k2_7.config import KimiK2Config + + torch.manual_seed(0) + cfg = KimiK2Config.reduced_marlin() + with torch.device("meta"): + block = KimiSparseMoeBlock(cfg) + block = block.to(torch.bfloat16) + block.to_empty(device=DEVICE) + + # Random router + shared-expert weights (float params only). + for p in block.parameters(): + if p.dtype.is_floating_point: + with torch.no_grad(): + p.copy_(torch.randn_like(p) * 0.1) + + # Synthetic quantized routed experts; keep the bf16 dequant for the reference. + E, H, I = cfg.n_routed_experts, cfg.hidden_size, cfg.moe_intermediate_size + w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16) + w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16) + w1_packed, w1_scale, w1_deq = _quantize_stack(w1) + w2_packed, w2_scale, w2_deq = _quantize_stack(w2) + with torch.no_grad(): + block.experts.gate_up_proj_packed.data = w1_packed + block.experts.gate_up_proj_scale.data = w1_scale + block.experts.down_proj_packed.data = w2_packed + block.experts.down_proj_scale.data = w2_scale + return cfg, block, (w1_deq, w2_deq) + + +def test_marlin_block_matches_bf16_reference(): + from mstar.utils.fused_moe.runner import fused_experts + + cfg, block, (w1_deq, w2_deq) = _build_block() + H = cfg.hidden_size + x = (torch.randn(5, H, device=DEVICE) * 0.5).to(torch.bfloat16) + + # bf16 reference (before the Marlin repack frees the packed params): router + + # bf16 fused experts on the dequant + the shared expert. + with torch.no_grad(): + topk_w, topk_ids = block.gate(x) + routed_ref = fused_experts(x, w1_deq, w2_deq, topk_w.to(x.dtype), topk_ids) + ref = (routed_ref + block.shared_expert(x)).view(x.shape) + + # Resolve backend + repack to Marlin, then run the block forward. + block.process_weights_after_loading(torch.device(DEVICE)) + assert block._use_marlin, "reduced_marlin config should select the Marlin backend" + # Source packed params are freed after the repack. + assert block.experts.gate_up_proj_packed.numel() == 0 + + with torch.no_grad(): + out = block(x) + + assert out.shape == x.shape and out.dtype == torch.bfloat16 + cos = torch.nn.functional.cosine_similarity( + out.flatten().float(), ref.flatten().float(), dim=0 + ).item() + rel_l2 = ((out - ref).float().norm() / ref.float().norm()).item() + assert cos > 0.999, f"cosine vs bf16 reference too low: {cos}" + assert rel_l2 < 0.02, f"relative-L2 vs bf16 reference too high: {rel_l2}" + + +def test_forced_marlin_raises_on_illegal_shapes(): + """``quant_kernel='marlin'`` must fail loudly when the shapes are Marlin-illegal + (rather than silently downgrading) — reduced_quantized_inkernel has + moe_intermediate_size=64, which violates Marlin's k%128 on the down GEMM.""" + from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock + from mstar.model.kimi_k2_7.config import KimiK2Config + + cfg = KimiK2Config.reduced_quantized_inkernel() + cfg.quant_kernel = "marlin" + with torch.device("meta"): + block = KimiSparseMoeBlock(cfg) + block = block.to(torch.bfloat16) + block.to_empty(device=DEVICE) + with pytest.raises(RuntimeError, match="Marlin is ineligible"): + block.process_weights_after_loading(torch.device(DEVICE)) + + +def test_triton_backend_still_selected_when_forced(): + """``quant_kernel='triton'`` keeps the packed Triton path (no Marlin repack).""" + from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock + from mstar.model.kimi_k2_7.config import KimiK2Config + + cfg = KimiK2Config.reduced_marlin() + cfg.quant_kernel = "triton" + with torch.device("meta"): + block = KimiSparseMoeBlock(cfg) + block = block.to(torch.bfloat16) + block.to_empty(device=DEVICE) + block.process_weights_after_loading(torch.device(DEVICE)) + assert not block._use_marlin + # Packed params are retained for the Triton kernel (not freed). + assert block.experts.gate_up_proj_packed.numel() > 0 diff --git a/test/integration/test_marlin_kernels.py b/test/integration/test_marlin_kernels.py new file mode 100644 index 000000000..0e9614bf1 --- /dev/null +++ b/test/integration/test_marlin_kernels.py @@ -0,0 +1,134 @@ +"""GPU golden for the vendored Marlin W4A16 kernels (utils/marlin). + +Validates the kernel layer in isolation, below any model wiring: + + 1. the JIT extension builds and registers ``_mstar_marlin_C`` ops; + 2. ``gptq_marlin_repack`` runs and is deterministic (repack layout is stable); + 3. the full routed-expert path (``MarlinMoEMethod`` = per-expert Marlin repack + + ``fused_marlin_moe``) matches the bf16 fused-expert GEMM on the *same* + dequantized weights, AND the existing Triton W4A16 path on the *same* packed + weights. + +Marlin dequantizes the identical INT4 nibbles as the Triton path but accumulates +in a different tile/reduce order (and folds fp32 combine weights), so agreement is +close-but-not-bit-exact. The meaningful gate is a cosine-similarity floor plus a +relative-L2 bound — an elementwise ``atol`` would flag the handful of bf16-accumulate +outliers on large down-projection sums, not a kernel bug. + +Run: pytest test/integration/test_marlin_kernels.py -v +""" +import pytest +import torch + +from mstar.model.kimi_k2_7._testing import fake_quantize_weight + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 0), + reason="Marlin W4A16 kernels need a CUDA GPU with sm80+ (Ampere/Hopper)", +) + +DEVICE = "cuda" +GROUP_SIZE = 32 +PACK_FACTOR = 8 + + +def _quantize_stack(weight): + """Fake-quantize a stacked ``(E, N, K)`` weight -> packed/scale/deq (bf16 scale).""" + E, N, K = weight.shape + packed = torch.empty((E, N, K // PACK_FACTOR), dtype=torch.int32, device=DEVICE) + scale = torch.empty((E, N, K // GROUP_SIZE), dtype=torch.bfloat16, device=DEVICE) + deq = torch.empty((E, N, K), dtype=torch.bfloat16, device=DEVICE) + for e in range(E): + p, s, d = fake_quantize_weight( + weight[e], num_bits=4, group_size=GROUP_SIZE, symmetric=True, + scale_dtype=torch.bfloat16, + ) + packed[e], scale[e], deq[e] = p.to(DEVICE), s.to(DEVICE), d.to(DEVICE) + return packed, scale, deq + + +def _random_topk(num_tokens, E, top_k): + logits = torch.randn(num_tokens, E, device=DEVICE) + weights, ids = torch.topk(logits.softmax(-1), top_k, dim=-1) + weights = weights / weights.sum(-1, keepdim=True) + return weights.to(torch.bfloat16), ids + + +def _rel_l2(a, b): + return ((a - b).float().norm() / b.float().norm()).item() + + +def test_marlin_builds_and_registers(): + from mstar.utils.marlin import is_marlin_available + + assert is_marlin_available(), "the vendored Marlin CUDA extension failed to build" + assert hasattr(torch.ops._mstar_marlin_C, "gptq_marlin_repack") + assert hasattr(torch.ops._mstar_marlin_C, "moe_wna16_marlin_gemm") + + +def test_repack_shape_and_deterministic(): + from mstar.utils.marlin import is_marlin_available, ops + + assert is_marlin_available() + size_k, size_n = 256, 128 # k%16==0, n%64==0 + b = torch.randint( + -(2**31), 2**31 - 1, (size_k // PACK_FACTOR, size_n), dtype=torch.int32, device=DEVICE + ) + out = ops.gptq_marlin_repack(b, size_k, size_n, num_bits=4) + assert out.shape == (size_k // 16, size_n * 16 // PACK_FACTOR) + assert torch.equal(out, ops.gptq_marlin_repack(b, size_k, size_n, num_bits=4)) + + +@pytest.mark.parametrize("num_tokens", [8, 3, 1]) # M > E, M <= E, single-token decode +def test_marlin_moe_matches_bf16_and_triton(num_tokens): + from mstar.model.components.quantization import MarlinMoEMethod + from mstar.utils.fused_moe.runner import fused_experts + + torch.manual_seed(0) + # Marlin-legal shapes: hidden % 128, moe_inter % 128, 2*inter % 64, hidden % 64. + E, H, I, top_k = 4, 256, 256, 2 + w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16) + w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16) + w1_packed, w1_scale, w1_deq = _quantize_stack(w1) + w2_packed, w2_scale, w2_deq = _quantize_stack(w2) + + x = (torch.randn(num_tokens, H, device=DEVICE) * 0.5).to(torch.bfloat16) + topk_weights, topk_ids = _random_topk(num_tokens, E, top_k) + + method = MarlinMoEMethod(num_bits=4, group_size=GROUP_SIZE) + method.prepare(w1_packed, w1_scale, w2_packed, w2_scale, torch.device(DEVICE)) + out_marlin = method.apply(x, topk_weights, topk_ids) + + out_bf16 = fused_experts(x, w1_deq, w2_deq, topk_weights, topk_ids) # ground truth + out_triton = fused_experts( # same packed nibbles, Triton W4A16 kernel + x, w1_packed, w2_packed, topk_weights, topk_ids, + w1_scale=w1_scale, w2_scale=w2_scale, group_size=GROUP_SIZE, pack_factor=PACK_FACTOR, + ) + + assert out_marlin.shape == (num_tokens, H) and out_marlin.dtype == torch.bfloat16 + cos = torch.nn.functional.cosine_similarity( + out_marlin.flatten().float(), out_bf16.flatten().float(), dim=0 + ).item() + assert cos > 0.999, f"cosine vs bf16 too low: {cos}" + assert _rel_l2(out_marlin, out_bf16) < 0.02, "relative-L2 vs bf16 too high" + assert _rel_l2(out_marlin, out_triton) < 0.02, "relative-L2 vs Triton W4A16 too high" + + +def test_marlin_moe_reduce_results_false_shape(): + """``reduce_results=False`` returns the per-slot (tokens, top_k, hidden) tensor + the TP path all-reduces before folding — exercise it on the Marlin path.""" + from mstar.model.components.quantization import MarlinMoEMethod + + torch.manual_seed(1) + E, H, I, top_k, num_tokens = 4, 256, 256, 2, 6 + w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16) + w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16) + w1_packed, w1_scale, _ = _quantize_stack(w1) + w2_packed, w2_scale, _ = _quantize_stack(w2) + x = (torch.randn(num_tokens, H, device=DEVICE) * 0.5).to(torch.bfloat16) + topk_weights, topk_ids = _random_topk(num_tokens, E, top_k) + + method = MarlinMoEMethod(num_bits=4, group_size=GROUP_SIZE) + method.prepare(w1_packed, w1_scale, w2_packed, w2_scale, torch.device(DEVICE)) + got = method.apply(x, topk_weights, topk_ids, reduce_results=False) + assert got.shape == (num_tokens, top_k, H) From 1eae55f43ab182350a6d24ffcd7cf4cdabd81d1c Mon Sep 17 00:00:00 2001 From: Garv Ghai Date: Sun, 26 Jul 2026 21:47:27 +0000 Subject: [PATCH 5/9] Kimi-K2.7: weight-absorbed MLA (default) + FlashInfer MLA kernel + CUDA-graph capture MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fold kv_b_proj into Q/O (w_kc/w_vc) + fuse q_a_proj/kv_a_proj_with_mqa, run MQA over a compressed-latent paged cache (~57× KV shrink). New MlaAbsorbCacheManager with a FlashInfer MLA kernel fast path (real dims/sm90) gated by a probe, SDPA fallback elsewhere; FlashInferMLAWrapper + _create_persistent_wrappers branch make absorbed decode CUDA-graph-capturable. Default via mla_absorb=True (reduced() pins naive as the parity reference). Composes with the Marlin MoE post-load walker. Validated e2e: real 1T serves at TP8 (absorbed kernel + Marlin + capture, coherent output); logits match vLLM up to bf16 near-ties. --- mstar/engine/cache_manager.py | 327 +++++++++++++++- mstar/engine/cuda_graph_runner.py | 36 +- mstar/engine/kv_cache_engine.py | 22 +- mstar/engine/kv_store.py | 12 + mstar/model/kimi_k2_7/components/attention.py | 159 +++++++- mstar/model/kimi_k2_7/config.py | 27 ++ mstar/model/kimi_k2_7/kimi_model.py | 36 +- mstar/utils/flashinfer_utils.py | 216 +++++++++++ .../test_kimi_mla_absorb_forward.py | 158 ++++++++ .../test_kimi_mla_absorb_kernel.py | 352 ++++++++++++++++++ .../test_kimi_mla_absorb_marlin_merge.py | 119 ++++++ .../integration/test_kimi_mla_absorb_paged.py | 182 +++++++++ .../integration/test_kimi_mla_absorb_serve.py | 262 +++++++++++++ test/modular/test_kimi_mla_absorb.py | 218 +++++++++++ 14 files changed, 2105 insertions(+), 21 deletions(-) create mode 100644 test/integration/test_kimi_mla_absorb_forward.py create mode 100644 test/integration/test_kimi_mla_absorb_kernel.py create mode 100644 test/integration/test_kimi_mla_absorb_marlin_merge.py create mode 100644 test/integration/test_kimi_mla_absorb_paged.py create mode 100644 test/integration/test_kimi_mla_absorb_serve.py create mode 100644 test/modular/test_kimi_mla_absorb.py diff --git a/mstar/engine/cache_manager.py b/mstar/engine/cache_manager.py index 557083e24..33d66a978 100644 --- a/mstar/engine/cache_manager.py +++ b/mstar/engine/cache_manager.py @@ -12,7 +12,11 @@ KVRequestState, PagedAllocationManager, ) -from mstar.utils.flashinfer_utils import FlashInferDecodeWrapper, FlashInferPrefillWrapper +from mstar.utils.flashinfer_utils import ( + FlashInferDecodeWrapper, + FlashInferMLAWrapper, + FlashInferPrefillWrapper, +) logger = logging.getLogger(__name__) @@ -98,7 +102,7 @@ class _PlanState: actually consumes this — the model's inner ``advance_seq_lens(pos_id_ns=...)`` runs at capture time only and is not replayed. """ - wrapper: FlashInferPrefillWrapper | FlashInferDecodeWrapper | None = None + wrapper: FlashInferPrefillWrapper | FlashInferDecodeWrapper | FlashInferMLAWrapper | None = None pos_ids: torch.Tensor | None = None seq_lens: list[int] | None = None write_store: bool = True @@ -123,6 +127,12 @@ class _PlanState: # segment over its contiguous frozen prefix. None on paged plans, which # keep the FlashInfer path. See DenseGenCacheManager._build_dense_gen_plan. dense_gen: dict | None = None + # Set when MlaAbsorbCacheManager planned this label: the compressed-latent + # scatter indices (token_to_page/token_to_cache) plus a per-request list of + # (q_start, seq_len, total_len, page_indices) that run_attention_mla uses to + # gather each request's full cached latent and run its causal SDPA. None on + # the standard FlashInfer/dense plans. See MlaAbsorbCacheManager.plan_attention. + mla: dict | None = None class WorkspaceBufferManager: @@ -1488,10 +1498,323 @@ def _run_dense_gen( return out[0] if isinstance(out, tuple) else out +@functools.cache +def _mla_kernel_available(ckv: int, kpe: int, sm_major: int) -> bool: + """Whether the FlashInfer MLA kernel fast path can serve these latent dims. + + Gated conservatively: ``flashinfer.mla.BatchMLAPagedAttentionWrapper`` is + hard-locked to the real Kimi dims (ckv=512, kpe=64). Off-dim calls trigger an + *uncatchable* illegal memory access (it corrupts the CUDA context — not a + catchable exception), so this decision MUST be made BEFORE any kernel + construction/call. Requires the flashinfer MLA module, exactly ckv=512/kpe=64, + and a Hopper (sm90) GPU (``backend="auto"`` -> fa3). Everything else — reduced + configs, pre-sm90, Blackwell (sm100, which wants the trtllm MLA path), or a + build without flashinfer — returns False and the manager uses the all-dims + SDPA fallback. + """ + if not (ckv == 512 and kpe == 64): + return False + if sm_major != 9: + return False + try: + import flashinfer.mla # noqa: F401 + except Exception: # noqa: BLE001 + return False + return True + + +class MlaAbsorbCacheManager(FlashInferCacheManager): + """Compressed-latent (weight-absorbed) MLA attention backend. + + Serves DeepSeek/Kimi MLA over a *latent* paged cache: instead of the + standard 6D ``[layers, pages, 2, page_size, kv_heads, head_dim]`` K/V cache, + the KV cache is the 4D latent tensor + ``[layers, pages, page_size, latent_width]`` allocated by + ``KVCacheEngine.load_model`` when ``attention_backend == "mla_absorb"`` + (``latent_width = kv_lora_rank + qk_rope_head_dim``). One compressed latent + vector ``cat([kv_c, k_pe])`` is stored per token (MQA: a single shared latent + "head"), which is ~L/head_dim smaller than materialized K/V. + + The standard FlashInfer paged wrappers assume the 6D layout, so this backend + does not build one. When the FlashInfer **MLA** kernel is available at the real + Kimi dims (see ``_mla_kernel_available``), ``plan_attention`` builds a + ``FlashInferMLAWrapper`` (the dedicated ckv=512/kpe=64 kernel) and + ``run_attention_mla`` scatters the new latents then runs it over strided + ``ckv``/``kpe`` views of the combined cache — this is the production fast path + and is CUDA-graph capturable. Otherwise (reduced configs, pre-sm90, no + flashinfer) it falls back to an all-dims **SDPA** path: ``plan_attention`` + records the per-token write (page, offset) indices — the same index math the + FlashInfer prefill wrapper's ``plan`` computes — plus the per-request gather + layout; ``run_attention_mla`` scatters the new latents, gathers each request's + full cached latent, and runs a causal SDPA in which ``value`` is the first + ``L`` dims of the same latent that forms ``key`` (weight absorption folds + ``k_nope``/``v`` into the query/output projections, so the cache only holds + ``kv_c`` + rope). The SDPA path is eager-only. + """ + + def plan_attention( + self, + seq_lens: list[int] | None = None, + dtype: torch.dtype | None = None, + is_causal=True, + write_store: bool=True, + label: str | None = None, + **kwargs, + ): + """Allocate pages and record the latent attention plan. + + Allocates enough pages for ``seq_len + new_tokens`` per request, then plans + one of two paths (chosen by ``_mla_kernel_available``): + + - **Kernel fast path** (real dims + sm90 + flashinfer): build the MLA index + tensors (qo_indptr / kv_indptr / kv_indices / kv_len_arr) — the same + batch-level tensors ``FlashInferCacheManager._plan_attention_impl`` builds + — and plan a persistent (CUDA-graph) or fresh (eager) ``FlashInferMLAWrapper`` + (stored on ``ps.wrapper``). ``ps.mla`` is cleared; the wrapper owns the + scatter indices. + - **SDPA fallback** (all dims, eager): record, for every new token, the + (page, within-page offset) it will be written to — computed exactly as + ``FlashInferPrefillWrapper.plan`` does (absolute position ``g`` maps to + page ``page_indices[g // page_size]`` at offset ``g % page_size``) — plus + the per-request gather layout, stashed on ``ps.mla``. ``ps.wrapper`` stays + None. + + Planning hints for other backends (``**kwargs``) are ignored. + """ + self._batched_cfg_info = None + + effective_label = label if label is not None else self._active_label() + # This backend always re-plans (it does not implement the plan-overlap + # short-circuit); clear any pre-plan marker so it never causes a stale skip. + # The re-plan writes the (per-slot) wrapper's static buffers correctly under + # capture — only the overlap perf optimization is forgone (a follow-up). + self._pre_planned_labels.discard(effective_label) + if effective_label not in self._plan_states: + self._plan_states[effective_label] = _PlanState() + ps = self._plan_states[effective_label] + + cfg = self.kv_cache_config + page_size = cfg.page_size + ckv = cfg.mla_ckv_dim + kpe = (cfg.head_dim - ckv) if ckv is not None else None + sm_major = torch.cuda.get_device_capability(self.device)[0] + use_kernel = ckv is not None and _mla_kernel_available(ckv, kpe, sm_major) + + if use_kernel: + # ---- Kernel fast path: build batched MLA index tensors + plan wrapper. + qo_indptr_list = [0] + kv_indptr_list = [0] + all_page_indices: list[int] = [] + kv_len_list: list[int] = [] + for i, rid in enumerate(self.request_ids): + state = self._get_state(rid, effective_label) + sl = seq_lens[i] + total_len = state.seq_len + sl + self.alloc_manager.alloc(rid, label=effective_label, seq_len=total_len) + page_indices = state.page_indices + qo_indptr_list.append(qo_indptr_list[-1] + sl) + all_page_indices.extend(page_indices) + kv_indptr_list.append(kv_indptr_list[-1] + len(page_indices)) + kv_len_list.append(total_len) + + qo_indptr = torch.tensor(qo_indptr_list, dtype=torch.int32, device=self.device) + kv_indptr = torch.tensor(kv_indptr_list, dtype=torch.int32, device=self.device) + kv_indices = torch.tensor(all_page_indices, dtype=torch.int32, device=self.device) + kv_len_arr = torch.tensor(kv_len_list, dtype=torch.int32, device=self.device) + + if dtype is None: + dtype = self.kv_cache.dtype + if ps.wrapper is None: + # Eager mode: fresh wrapper each forward (the manager is rebuilt per + # forward). CUDA-graph mode passes a persistent wrapper in ps.wrapper. + ps.wrapper = FlashInferMLAWrapper( + workspace_buffer=self.buffer_manager.get(effective_label), + num_heads=cfg.num_qo_heads, + head_dim_ckv=ckv, + head_dim_kpe=kpe, + page_size=page_size, + sm_scale=cfg.softmax_scale, + device=self.device, + enable_nvtx=self.enable_nvtx, + ) + ps.wrapper.plan( + qo_indptr=qo_indptr, + kv_indptr=kv_indptr, + kv_indices=kv_indices, + kv_len_arr=kv_len_arr, + causal=is_causal, + dtype=dtype, + ) + ps.mla = None + else: + # ---- SDPA fallback: per-token scatter + per-request gather layout. + token_to_page: list[int] = [] + token_to_cache: list[int] = [] + requests: list[dict] = [] + q_start = 0 + for i, rid in enumerate(self.request_ids): + state = self._get_state(rid, effective_label) + sl = seq_lens[i] + old_len = state.seq_len + total_len = old_len + sl + + self.alloc_manager.alloc( + rid, label=effective_label, seq_len=total_len + ) + page_indices = state.page_indices + + # WRITE indices for the sl new tokens: token j lands at absolute + # position g = old_len + j -> page page_indices[g // page_size], + # offset g % page_size. (Same math as the FlashInfer plan's + # token_to_page / token_to_cache, specialized to per-request order.) + for j in range(sl): + g = old_len + j + token_to_page.append(page_indices[g // page_size]) + token_to_cache.append(g % page_size) + + requests.append({ + "q_start": q_start, + "seq_len": sl, + "total_len": total_len, + "page_indices": torch.tensor( + page_indices, dtype=torch.long, device=self.device + ), + }) + q_start += sl + + ps.mla = { + "token_to_page": torch.tensor( + token_to_page, dtype=torch.long, device=self.device + ), + "token_to_cache": torch.tensor( + token_to_cache, dtype=torch.long, device=self.device + ), + "requests": requests, + } + + # Record like the base plan does so advance_seq_lens / flush_to_store + # see this label's per-request new-token counts and store policy. + ps.seq_lens = seq_lens + ps.write_store = write_store + ps.dense_gen = None + + @torch.compiler.disable + def run_attention_mla( + self, + q_nope: torch.Tensor, + q_pe: torch.Tensor, + kv_c: torch.Tensor, + k_pe: torch.Tensor, + layer_idx: int | None = None, + ) -> torch.Tensor: + """Compressed-latent MLA attention over the paged latent cache. + + Args: + q_nope: [T, H, L] query, no-rope part (L = kv_lora_rank). + q_pe: [T, H, Drope] query, rope part. + kv_c: [T, 1, L] compressed KV latent (single MQA head). + k_pe: [T, 1, Drope] rope key (single MQA head). + layer_idx: transformer layer; defaults to self.layer_idx. + Returns: + [T, H, L] attention output (the ``value`` = ``kv_c`` slice width). + + Writes ``cat([kv_c, k_pe])`` (width L+Drope) as one latent per token into + this layer's paged latent cache at the planned (page, offset) locations, + then attends. When ``plan_attention`` selected the FlashInfer MLA kernel + (``ps.wrapper`` set), the kernel runs over strided ``ckv``/``kpe`` views of + the combined cache; otherwise a causal SDPA gathers each request's full + cached latent (query ``cat([q_nope, q_pe])``, key = the full latent, value = + its first L dims) at ``kv_cache_config.softmax_scale``. + """ + if layer_idx is None: + layer_idx = self.layer_idx + + label = self._active_label() + ps = self._plan_states[label] + assert self.kv_cache is not None + + latent_cache = self.kv_cache[layer_idx] # [max_pages, page_size, L+Drope] + latent = torch.cat([kv_c, k_pe], dim=-1).squeeze(1) # [T, L+Drope] + + if ps.wrapper is not None: + # ---- FlashInfer MLA kernel fast path (CUDA-graph capturable). + ps.wrapper.set_latent(latent_cache, latent) + L = q_nope.shape[-1] # ckv width (post-w_kc absorption) + ckv_cache = latent_cache[..., :L] + kpe_cache = latent_cache[..., L:] + return ps.wrapper.run(q_nope, q_pe, ckv_cache, kpe_cache).to(q_nope.dtype) + + # ---- SDPA fallback. + mla = ps.mla + assert mla is not None + + # Scatter: one latent vector per new token into (page, offset). + latent_cache[mla["token_to_page"], mla["token_to_cache"]] = latent.to( + latent_cache.dtype + ) + + T, H, L = q_nope.shape + scale = self.kv_cache_config.softmax_scale + query_all = torch.cat([q_nope, q_pe], dim=-1) # [T, H, L+Drope] + out = torch.empty(T, H, L, dtype=q_nope.dtype, device=q_nope.device) + + for req in mla["requests"]: + q_start = req["q_start"] + sl = req["seq_len"] + total_len = req["total_len"] + + # Gather this request's full cached latent [total_len, L+Drope] + # (mirror the DenseGenCacheManager page-gather). + gathered = latent_cache[req["page_indices"]].reshape( + -1, latent_cache.shape[-1] + )[:total_len] + key = gathered # [total_len, L+Drope] + value = gathered[:, :L] # [total_len, L] + + q_req = query_all[q_start:q_start + sl] # [sl, H, L+Drope] + out[q_start:q_start + sl] = self._sdpa_mla( + q_req, key, value, old_len=total_len - sl, scale=scale + ) + return out + + @staticmethod + def _sdpa_mla( + q: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + old_len: int, + scale: float, + ) -> torch.Tensor: + """Causal SDPA for one request's MLA step. + + ``q`` is [sl, H, D] (D = L+Drope); ``key`` [total, D] and ``value`` + [total, L] are the single latent "head" broadcast over the H query heads. + Query token j is at absolute position ``old_len + j`` and attends to + cached positions ``0 .. old_len + j`` (causal, includes itself). Handles + both prefill (old_len=0) and a decode step (sl=1, old_len=total-1). + """ + sl = q.shape[0] + total = key.shape[0] + qt = q.transpose(0, 1).float() # [H, sl, D] + scores = torch.einsum("hqd,kd->hqk", qt, key.float()) * scale # [H, sl, total] + q_pos = old_len + torch.arange(sl, device=q.device) + k_pos = torch.arange(total, device=q.device) + mask = torch.where( + k_pos[None, :] <= q_pos[:, None], + 0.0, + torch.tensor(float("-inf"), device=q.device), + ) + scores = scores + mask + attn = scores.softmax(-1) + out = torch.einsum("hqk,kd->hqd", attn, value.float()) # [H, sl, L] + return out.transpose(0, 1).to(q.dtype) # [sl, H, L] + + # Backend registry: KVCacheConfig.attention_backend names one of these. ATTENTION_BACKENDS: dict[str, type[BatchedCacheManager]] = { "flashinfer": FlashInferCacheManager, "dense_gen": DenseGenCacheManager, + "mla_absorb": MlaAbsorbCacheManager, } diff --git a/mstar/engine/cuda_graph_runner.py b/mstar/engine/cuda_graph_runner.py index a9252e3b6..6da5db6d2 100644 --- a/mstar/engine/cuda_graph_runner.py +++ b/mstar/engine/cuda_graph_runner.py @@ -326,9 +326,10 @@ def _create_persistent_wrappers( can run on plan_stream concurrently with replay(slot 0) on default_stream without racing on the wrapper's persistent state. """ - from mstar.engine.cache_manager import _PlanState + from mstar.engine.cache_manager import _mla_kernel_available, _PlanState from mstar.utils.flashinfer_utils import ( FlashInferDecodeWrapper, + FlashInferMLAWrapper, FlashInferPrefillWrapper, ) @@ -336,6 +337,22 @@ def _create_persistent_wrappers( cfg = self.kv_cache_config + # Compressed-latent MLA fast path: when the backend is "mla_absorb" and the + # FlashInfer MLA kernel is available at these dims (real Kimi dims on sm90), + # capture uses a persistent FlashInferMLAWrapper for BOTH decode and prefill + # (its run() serves both). This is the only capturable absorbed path — the + # SDPA fallback is eager-only, so reduced-dims / non-kernel absorbed serving + # runs eager (no capture). See MlaAbsorbCacheManager. + use_mla_kernel = ( + cfg.attention_backend == "mla_absorb" + and cfg.mla_ckv_dim is not None + and _mla_kernel_available( + cfg.mla_ckv_dim, + cfg.head_dim - cfg.mla_ckv_dim, + torch.cuda.get_device_capability(self.device)[0], + ) + ) + # Allocate workspace buffer for CUDA graph wrappers. # Each (label, slot) gets its own workspace — slots must NOT share # workspace because plan() writes scheduling state there and the @@ -344,7 +361,22 @@ def _create_persistent_wrappers( plan_states = {} for label in config.labels: ws_label = f"{label}_cugraph_slot{slot_idx}" - if is_decode: + if use_mla_kernel: + wrapper = FlashInferMLAWrapper( + workspace_buffer=self.buffer_manager.get(ws_label), + num_heads=cfg.num_qo_heads, + head_dim_ckv=cfg.mla_ckv_dim, + head_dim_kpe=cfg.head_dim - cfg.mla_ckv_dim, + page_size=cfg.page_size, + sm_scale=cfg.softmax_scale, + batch_size=bs, + max_num_pages=cfg.max_num_pages, + max_total_tokens=total_tokens, + device=self.device, + use_cuda_graph=True, + enable_nvtx=self.enable_nvtx, + ) + elif is_decode: wrapper = FlashInferDecodeWrapper( workspace_buffer=self.buffer_manager.get(ws_label), num_qo_heads=cfg.num_qo_heads, diff --git a/mstar/engine/kv_cache_engine.py b/mstar/engine/kv_cache_engine.py index df85289b5..7671d30b3 100644 --- a/mstar/engine/kv_cache_engine.py +++ b/mstar/engine/kv_cache_engine.py @@ -228,11 +228,23 @@ def load_model( ) num_kv_heads = cfg.num_kv_heads - kv_cache = torch.zeros( - num_layers, max_num_pages, 2, - page_size, num_kv_heads, head_dim, - dtype=kv_cache_type, device=device, - ).contiguous() + if cfg.attention_backend == "mla_absorb": + # Compressed-latent MLA cache: one latent vector per token of + # width head_dim (= kv_lora_rank + qk_rope_head_dim). Drop the + # 2-wide K/V axis (a single latent, not a K/V pair) and the + # num_kv_heads axis (MQA — one shared latent head), giving a 4D + # [num_layers, max_pages, page_size, latent_width] cache that + # MlaAbsorbCacheManager scatters into / gathers from. + kv_cache = torch.zeros( + num_layers, max_num_pages, page_size, head_dim, + dtype=kv_cache_type, device=device, + ).contiguous() + else: + kv_cache = torch.zeros( + num_layers, max_num_pages, 2, + page_size, num_kv_heads, head_dim, + dtype=kv_cache_type, device=device, + ).contiguous() cpu_page_pool = None if cfg.cpu_offload_pages > 0: diff --git a/mstar/engine/kv_store.py b/mstar/engine/kv_store.py index 3db7606ad..465431f25 100644 --- a/mstar/engine/kv_store.py +++ b/mstar/engine/kv_store.py @@ -122,6 +122,18 @@ class KVCacheConfig: # FA3 on Hopper; models can pin ``fa2`` when their deployment toolchain # cannot compile the Hopper JIT kernels. flashinfer_backend: str = "auto" + # Softmax scale for the compressed-latent MLA backend ("mla_absorb"). MLA's + # intended scale is qk_head_dim**-0.5 * mscale**2, which differs from the + # 1/sqrt(head_dim) a standard kernel would apply over the latent width, so + # the model passes the correct value here and MlaAbsorbCacheManager reads it + # in run_attention_mla. None for the standard paged backends (unused). + softmax_scale: float | None = None + # For "mla_absorb": the compressed-KV latent width (kv_lora_rank), i.e. the + # ``ckv`` half of the combined latent ``head_dim = ckv + kpe``. The FlashInfer + # MLA kernel fast path needs this split at plan time to pass head_dim_ckv / + # head_dim_kpe. None for the standard paged backends (unused; SDPA fallback + # derives the split from the query shapes at run time). + mla_ckv_dim: int | None = None def __post_init__(self): if self.num_qo_heads is None: diff --git a/mstar/model/kimi_k2_7/components/attention.py b/mstar/model/kimi_k2_7/components/attention.py index fc861d17a..d611adf57 100644 --- a/mstar/model/kimi_k2_7/components/attention.py +++ b/mstar/model/kimi_k2_7/components/attention.py @@ -1,11 +1,26 @@ -"""Kimi-K2.7 / DeepSeek-V3 MLA attention — naive / materialized path. +"""Kimi-K2.7 / DeepSeek-V3 MLA attention — naive path + weight-absorbed path. -MLA compresses q and k/v through low-rank latents, then (in the naive path) -projects the latent back up to full per-head K/V and runs ordinary attention. -This avoids the weight-absorbed path (``W_UK``/``W_UV``) and its bespoke kernel — -throughput caveat noted — so it drops straight onto mstar's paged -``run_attention`` ``[tokens, heads, head_dim]`` interface, matching vLLM's -``DeepseekV2Attention`` (the non-absorbed class). +MLA compresses q and k/v through low-rank latents. Two forwards live here, picked +by ``config.mla_absorb`` (default ``True`` -> absorbed): + +* **weight-absorbed** (``mla_absorb=True``, DEFAULT): folds ``kv_b_proj``'s + up-projection into the Q path (``W_UK``) and the O path (``W_UV``) at load (plus + the ``fused_qkv_a_proj`` down-proj fusion), via + :meth:`KimiMLAAttention.process_weights_after_loading`, so attention runs as MQA + over the COMPRESSED latent (``kv_c | k_pe``, one KV head) via + ``cache_handle.run_attention_mla`` — a ~57x per-token KV shrink, numerically + identical to naive up to fp rounding. See :meth:`_forward_absorbed`. Served by + ``engine/cache_manager.py::MlaAbsorbCacheManager`` over a 4D latent paged cache. + That backend currently uses a torch SDPA-over-latent path (correct + memory-lean + but EAGER-ONLY, no CUDA-graph capture); the FlashInfer MLA kernel + CUDA-graph + capture for real-1T throughput is a follow-up. +* **naive / materialized** (``mla_absorb=False``): projects the latent back up to + full per-head K/V and runs ordinary attention, dropping straight onto mstar's + paged ``run_attention`` ``[tokens, heads, head_dim]`` interface (matching vLLM's + ``DeepseekV2Attention``, the non-absorbed class). Zero-pads q/k/v to + ``padded_head_dim`` and folds a softmax boost into q. The M4-golden parity + reference + opt-out fallback (production should keep this until the MLA kernel + lands). Per-token shape story (H heads, Dnope=qk_nope, Drope=qk_rope, Dqk=Dnope+Drope, Dv=v_head_dim, L=kv_lora_rank): @@ -28,8 +43,15 @@ axis — there is no separate KV-head group in the naive path), and the paged cache reports the matching per-rank count. -TODO: weight-absorbed MLA (native latent dims, no pad) and the ``fused_qkv_a_proj`` -weight fusion are not implemented yet. +The absorbed cache config (``num_kv_heads == 1``, ``head_dim == kv_lora_rank + +qk_rope_head_dim``, no pad) is reported by ``kimi_model.py::get_kv_cache_config`` +when ``mla_absorb`` is set. + +The ``fused_qkv_a_proj`` weight fusion (``q_a_proj`` + ``kv_a_proj_with_mqa`` -> one +GEMM) is applied in the ABSORBED path only: ``process_weights_after_loading`` concats +the two replicated latent down-projection weights into ``fused_qkv_a_proj_weight`` and +:meth:`_forward_absorbed` runs a single ``F.linear`` then splits. The naive ``forward`` +keeps the two separate calls unchanged. """ from __future__ import annotations @@ -132,6 +154,24 @@ def __init__(self, config: KimiK2Config, comm_group: CommGroup | None = None) -> self.softmax_scale_boost = ( mscale * mscale * math.sqrt(self.padded_head_dim / self.qk_head_dim) ) + # DeepSeek's intended softmax scale (uses the PRE-absorption qk_head_dim, + # not the latent width). The naive path reaches it by folding + # softmax_scale_boost into q; the absorbed path hands this to the latent + # cache backend (planned in Phase B; the reduced-config test mock reads it). + self.softmax_scale = self.qk_head_dim ** -0.5 * mscale * mscale + + # Weight-absorbed MLA (config.mla_absorb): W_UK/W_UV are split out of + # kv_b_proj post-load by process_weights_after_loading(). persistent=False: + # derived tensors, never part of the checkpoint / state_dict. + self.mla_absorb = config.mla_absorb + if self.mla_absorb: + self.register_buffer("w_kc", None, persistent=False) # (H_local, Dnope, L) + self.register_buffer("w_vc", None, persistent=False) # (H_local, Dv, L) + # The two replicated latent down-projections (q_a_proj + + # kv_a_proj_with_mqa) both read hidden_states; they fuse into one GEMM + # for the absorbed forward. Built post-load; q_a_proj/kv_a_proj_with_mqa + # remain the checkpoint load targets. persistent=False: derived tensor. + self.register_buffer("fused_qkv_a_proj_weight", None, persistent=False) # (q_lora+L+Drope, hidden) def forward( self, @@ -139,6 +179,8 @@ def forward( cache_handle: BatchedCacheManager, position_ids: torch.Tensor, ) -> torch.Tensor: + if self.mla_absorb: + return self._forward_absorbed(hidden_states, cache_handle, position_ids) num_tokens = hidden_states.shape[0] h = self.num_heads @@ -176,3 +218,102 @@ def forward( attn = cache_handle.run_attention(q=q, k=k, v=v) # (T, H, Dpad) attn = attn[..., : self.v_head_dim].reshape(num_tokens, h * self.v_head_dim) return self.o_proj(attn) + + def _forward_absorbed( + self, + hidden_states: torch.Tensor, + cache_handle: BatchedCacheManager, + position_ids: torch.Tensor, + ) -> torch.Tensor: + """Weight-absorbed MLA forward (see module docstring; ``config.mla_absorb``). + + kv_b_proj is folded into Q (``W_UK`` = ``w_kc``) and O (``W_UV`` = ``w_vc``) + so the KV latent stays COMPRESSED and attention is MQA over ``[kv_c | k_pe]`` + (one KV head). Math identity vs naive: ``q_nope · k_nope == + (q_nope @ W_UK) · kv_c`` and ``attn · v == (attn · kv_c) @ W_UV``. The + softmax scale (``self.softmax_scale``) is applied by the latent cache + backend (planned in Phase B; the reduced-config test mock reads it), so — + unlike naive — nothing is folded into q and nothing is padded. + """ + if self.w_kc is None or self.w_vc is None: + raise RuntimeError( + "mla_absorb forward requires process_weights_after_loading() to " + "have built w_kc/w_vc from kv_b_proj first" + ) + if self.fused_qkv_a_proj_weight is None: + raise RuntimeError( + "mla_absorb forward requires process_weights_after_loading() to " + "have built fused_qkv_a_proj_weight first" + ) + num_tokens = hidden_states.shape[0] + h = self.num_heads + + # --- fused latent down-projection (q_a_proj + kv_a_proj_with_mqa in one + # GEMM), then split into the q latent, the kv latent, and the rope key --- + fused = F.linear(hidden_states, self.fused_qkv_a_proj_weight) + q_c, kv_a, k_pe = fused.split( + [self.q_a_proj.out_features, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + # q_c / kv_a feed the FlashInfer RMSNorm kernel, which requires a + # 64-byte-aligned input pointer. These are mid-tensor split views: q_c is at + # offset 0 (always aligned); kv_a sits at q_lora_rank*dtype bytes (e.g. + # reduced dims -> 96, not 64-aligned). ``.contiguous()`` fixes the multi-token + # (prefill) case but is a NO-OP for a size-1 (decode, T=1) row — a [1,K] view + # is already "contiguous" — leaving kv_a at the misaligned offset. So force a + # fresh contiguous allocation for kv_a via ``clone``. (Real dims align by + # chance; this keeps reduced-config / decode correct too.) + q_c = q_c.contiguous() + kv_a = kv_a.clone(memory_format=torch.contiguous_format) + + # --- Q (norm -> up), split nope/rope (same as naive) --- + q = self.q_b_proj(self.q_a_layernorm(q_c)) + q = q.view(num_tokens, h, self.qk_head_dim) + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + # --- KV latent kept COMPRESSED (no kv_b up-projection); norm the kv_c slice --- + kv_c = self.kv_a_layernorm(kv_a).view(num_tokens, 1, self.kv_lora_rank) # (T,1,L) + k_pe = k_pe.view(num_tokens, 1, self.qk_rope_head_dim) # (T,1,Drope) shared MQA key + + # --- RoPE on the pe slices only (q_pe per head, k_pe single shared key) --- + q_pe, k_pe = self.rotary(position_ids, q_pe, k_pe) + + # --- absorb W_UK into q_nope: (T,H,Dnope) x (H,Dnope,L) -> (T,H,L) --- + q_nope = torch.einsum("thd,hdl->thl", q_nope, self.w_kc) + + # --- MQA over the latent: key=[kv_c|k_pe], value=kv_c -> (T,H,L) --- + attn_latent = cache_handle.run_attention_mla( + q_nope=q_nope, q_pe=q_pe, kv_c=kv_c, k_pe=k_pe) + + # --- absorb W_UV into the output: (T,H,L) x (H,Dv,L) -> (T,H,Dv) -> o_proj --- + out = torch.einsum("thl,hdl->thd", attn_latent, self.w_vc) + return self.o_proj(out.reshape(num_tokens, h * self.v_head_dim)) + + def process_weights_after_loading(self, device: torch.device | str | None = None) -> None: + """Build the absorbed projections from ``kv_b_proj`` (no-op unless mla_absorb). + + Splits ``kv_b_proj.weight`` ``[H_local*(Dnope+Dv), L]`` per local head into + ``w_kc = W_UK [H_local, Dnope, L]`` (absorbed into Q) and + ``w_vc = W_UV [H_local, Dv, L]`` (absorbed into O). Each rank's ``kv_b_proj`` + (ColumnParallelLinear) already holds only its local heads, so this is + TP-correct with no extra sharding — one path serves tp=1 and tp>1. + + Named to match the generic post-load walker protocol so the serve path picks + it up automatically (Phase B). Idempotent. ``device`` is accepted for + protocol compatibility but unused — ``kv_b_proj.weight`` is already resident. + """ + if not self.mla_absorb: + return + del device # protocol arg; kv_b_proj.weight already carries the right device + w = self.kv_b_proj.weight # (H_local*(Dnope+Dv), L) + h, d_nope, d_v, latent = ( + self.num_heads, self.qk_nope_head_dim, self.v_head_dim, self.kv_lora_rank) + w = w.view(h, d_nope + d_v, latent) + w_kc, w_vc = w.split([d_nope, d_v], dim=1) + self.w_kc = w_kc.contiguous() # (H_local, Dnope, L) + self.w_vc = w_vc.contiguous() # (H_local, Dv, L) + + # Fuse the two replicated latent down-projections into one GEMM weight for + # the absorbed forward: [q_a_proj ; kv_a_proj_with_mqa] -> a single + # (q_lora + kv_lora + Drope, hidden) matmul that splits back into the q + # latent, the kv latent, and the shared rope key. + self.fused_qkv_a_proj_weight = torch.cat( + [self.q_a_proj.weight, self.kv_a_proj_with_mqa.weight], dim=0).contiguous() diff --git a/mstar/model/kimi_k2_7/config.py b/mstar/model/kimi_k2_7/config.py index fc2b51941..323a5ec12 100644 --- a/mstar/model/kimi_k2_7/config.py +++ b/mstar/model/kimi_k2_7/config.py @@ -43,6 +43,26 @@ class KimiK2Config: qk_rope_head_dim: int = 64 v_head_dim: int = 128 + # -- MLA weight absorption (DEFAULT) ----------------------------------- + # ``True`` (default) => weight-absorbed MLA: ``kv_b_proj``'s up-projection is + # folded into the Q path (``W_UK``) and the O path (``W_UV``) at load (plus the + # ``fused_qkv_a_proj`` down-proj fusion), attention runs as MQA over the + # COMPRESSED latent via the ``mla_absorb`` cache backend, and the KV cache + # stores only the ``kv_lora_rank + qk_rope_head_dim`` latent (1 KV head) — a + # ~57x per-token cache shrink, numerically identical up to fp rounding. + # ``False`` => naive/materialized MLA (latent projected up to full per-head + # K/V, padded to ``padded_head_dim``, MHA cache): the M4-golden parity + # reference / opt-out fallback. + # + # PERF CAVEAT: the absorbed backend currently runs on a torch SDPA-over-latent + # path — correct + memory-lean but EAGER-ONLY (no CUDA-graph capture) and slow + # on the real 1T. The FlashInfer MLA kernel + CUDA-graph capture (production + # throughput) is a follow-up; until it lands, real large-scale serving should + # set ``mla_absorb=False`` (naive) or accept eager execution. See + # ``components/attention.py``, ``kimi_model.py::get_kv_cache_config``, + # ``engine/cache_manager.py::MlaAbsorbCacheManager``. + mla_absorb: bool = True + # -- Fine-grained MoE (sigmoid router, group-limited top-k, noaux_tc) -- n_routed_experts: int = 384 # from config.json n_shared_experts: int = 1 @@ -161,8 +181,15 @@ def reduced(cls) -> "KimiK2Config": reduced-config golden runs. Keeps the *shape* of Kimi (MLA split heads, grouped MoE, one dense layer) while being small enough to run without the 1T checkpoint. + + NOTE ``mla_absorb=False``: this fixture pins the NAIVE MLA path, the + M4-golden parity reference that the bulk of the reduced test suite + validates. The absorbed path (the production default) is exercised by the + dedicated ``test_kimi_mla_absorb*`` tests, which flip ``mla_absorb=True`` + on a reduced() instance explicitly. """ return cls( + mla_absorb=False, vocab_size=256, hidden_size=128, intermediate_size=256, diff --git a/mstar/model/kimi_k2_7/kimi_model.py b/mstar/model/kimi_k2_7/kimi_model.py index 331272944..7c00fe605 100644 --- a/mstar/model/kimi_k2_7/kimi_model.py +++ b/mstar/model/kimi_k2_7/kimi_model.py @@ -119,6 +119,32 @@ def tokenizer(self): # ------------------------------------------------------------------- def get_kv_cache_config(self) -> list[KVCacheConfig]: + if self.config.mla_absorb: + # Weight-absorbed MLA (the default): attention is MQA over the + # COMPRESSED latent (kv_b_proj folded into Q/O + the fused_qkv_a_proj + # down-proj), so the paged cache stores a single KV "head" of width + # ``kv_lora_rank + qk_rope_head_dim`` ([kv_c | k_pe]) per token — a ~57x + # shrink vs the naive padded MHA cache (real 2*64*256=32768 -> 512+64=576; + # reduced 2*4*64=512 -> 40). Served by ``MlaAbsorbCacheManager`` over the + # 4D latent cache. ``softmax_scale`` is DeepSeek's intended MLA scale + # ``qk_head_dim**-0.5 * mscale**2``: the absorbed forward folds nothing + # into q (unlike naive), so the backend applies this scale directly. + from mstar.model.kimi_k2_7.components.rope import yarn_get_mscale + rope = self.config.rope_scaling + mscale = yarn_get_mscale(rope["factor"], rope.get("mscale_all_dim", 0.0)) + softmax_scale = self.config.qk_head_dim ** -0.5 * mscale * mscale + return [KVCacheConfig( + num_layers=self.config.num_hidden_layers, + num_kv_heads=1, + head_dim=self.config.kv_lora_rank + self.config.qk_rope_head_dim, + max_seq_len=self.config.max_position_embeddings, + num_qo_heads=self.config.num_attention_heads, + attention_backend="mla_absorb", + softmax_scale=softmax_scale, + # The ckv/kpe split for the FlashInfer MLA kernel fast path + # (head_dim = ckv + kpe = kv_lora_rank + qk_rope_head_dim). + mla_ckv_dim=self.config.kv_lora_rank, + )] # Naive/materialized MLA (the first-pass port, per CLAUDE.md): the latent # is projected up to full per-head K/V and broadcast to every query head, # so from the paged cache's ``[tokens, heads, head_dim]`` point of view @@ -406,9 +432,13 @@ def _create_submodule( language_model = language_model.to(autocast_dtype) language_model.to_empty(device=device) load_weights(language_model, source, device=device) - # Post-load pass: let quantized submodules finalize their kernel layout on - # the real device (the routed-expert MoE block repacks its packed experts - # into Marlin layout + allocates a workspace). No-op for a plain bf16 build. + # Post-load pass: let any submodule finalize its weights on the real device + # now that the checkpoint is resident. The routed-expert MoE block repacks + # its packed experts into the Marlin (or Triton) kernel layout + allocates a + # workspace; under mla_absorb, KimiMLAAttention builds the absorbed w_kc/w_vc + # projections + the fused_qkv_a_proj weight. Generic + idempotent walker + # (calls the hook on every module that exposes one); no-op for a plain bf16 + # module without the hook. from mstar.model.components.quantization import process_weights_after_loading process_weights_after_loading(language_model, torch.device(device)) diff --git a/mstar/utils/flashinfer_utils.py b/mstar/utils/flashinfer_utils.py index 899bdee3c..a736ca6aa 100644 --- a/mstar/utils/flashinfer_utils.py +++ b/mstar/utils/flashinfer_utils.py @@ -487,3 +487,219 @@ def set_kv_cache( positions = self.kv_cache_locations[:n, 1] kv_cache_layer[pages, 0, positions] = k[:n].to(self.dtype) kv_cache_layer[pages, 1, positions] = v[:n].to(self.dtype) + + +class FlashInferMLAWrapper: + """Compressed-latent (weight-absorbed) MLA attention over a paged latent cache. + + Wraps ``flashinfer.mla.BatchMLAPagedAttentionWrapper`` (the DeepSeek/Kimi MLA + kernel) for the ``MlaAbsorbCacheManager`` fast path. Unlike the standard + Prefill/Decode wrappers this consumes a **4D latent cache** + ``[max_pages, page_size, head_dim_ckv + head_dim_kpe]`` (one shared MQA latent + per token) and takes the query pre-split into its no-rope / rope parts: + + run(q_nope[T, H, ckv], q_pe[T, H, kpe], + ckv_cache[pages, page_size, ckv], kpe_cache[pages, page_size, kpe]) + -> [T, H, ckv] + + ``ckv``/``kpe`` are passed as **strided views** of the combined latent cache + (``cache[..., :ckv]`` / ``cache[..., ckv:]``) — the kernel accepts them, so the + cache stays a single tensor. + + The kernel is hard-locked to the real Kimi dims (ckv=512, kpe=64): other dims + trigger an *uncatchable* illegal memory access, so the caller must gate on dims + before constructing this (see ``cache_manager._mla_kernel_available``). + + Mirrors the Prefill/Decode wrappers: ``plan`` computes the per-token latent + scatter indices (into static buffers under CUDA graph); ``set_latent`` scatters + ``cat([kv_c, k_pe])`` into the cache; ``run`` calls the kernel. CUDA graph mode + requires ``batch_size``/``max_num_pages`` and static index buffers so ``plan`` + updates values via ``.copy_()`` without reallocating. + """ + + def __init__( + self, + workspace_buffer: torch.Tensor, + *, + num_heads: int, + head_dim_ckv: int, + head_dim_kpe: int, + page_size: int, + sm_scale: float, + batch_size: int | None = None, + max_num_pages: int | None = None, + max_total_tokens: int | None = None, + device: torch.device = torch.device("cuda"), + use_cuda_graph: bool = False, + backend: str = "auto", + enable_nvtx: bool = False, + ): + self.device = device + self.use_cuda_graph = use_cuda_graph + self.enable_nvtx = enable_nvtx + self.num_heads = num_heads + self.head_dim_ckv = head_dim_ckv + self.head_dim_kpe = head_dim_kpe + self.page_size = page_size + self.sm_scale = sm_scale + self.batch_size = batch_size + self.max_total_tokens = max_total_tokens + self.dtype = None + self._total_tokens = 0 + + import flashinfer + + if self.use_cuda_graph: + assert batch_size is not None, "batch_size required for CUDA graph mode" + assert max_num_pages is not None, "max_num_pages required for CUDA graph mode" + assert max_total_tokens is not None, "max_total_tokens required for CUDA graph mode" + + # Static index buffers the kernel plan copies into (see the wrapper's + # __init__ docstring): stable addresses across graph replay. + self._qo_indptr_buf = torch.zeros( + batch_size + 1, dtype=torch.int32, device=device + ) + self._kv_indptr_buf = torch.zeros( + batch_size + 1, dtype=torch.int32, device=device + ) + self._kv_indices_buf = torch.zeros( + max_num_pages, dtype=torch.int32, device=device + ) + self._kv_len_arr_buf = torch.zeros( + batch_size, dtype=torch.int32, device=device + ) + + self.attn_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( + workspace_buffer, + use_cuda_graph=True, + qo_indptr=self._qo_indptr_buf, + kv_indptr=self._kv_indptr_buf, + kv_indices=self._kv_indices_buf, + kv_len_arr=self._kv_len_arr_buf, + backend=backend, + ) + + # Static buffers for the vectorized latent scatter (page, offset). + self.token_to_page = torch.zeros( + max_total_tokens, dtype=torch.long, device=device + ) + self.token_to_cache = torch.zeros( + max_total_tokens, dtype=torch.long, device=device + ) + else: + self.attn_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( + workspace_buffer, backend=backend, + ) + self.token_to_page = None + self.token_to_cache = None + + @torch.compiler.disable + def plan( + self, + qo_indptr: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + kv_len_arr: torch.Tensor, + *, + causal: bool = True, + dtype: torch.dtype = torch.bfloat16, + ): + """Plan the MLA kernel and compute the per-token latent scatter indices. + + Args (all int32, on ``self.device``): + qo_indptr: [n_req + 1] cumulative NEW query tokens per request. + kv_indptr: [n_req + 1] cumulative pages per request. + kv_indices: [total_pages] flattened page indices (per-request order). + kv_len_arr: [n_req] total cached length per request AFTER this append. + + In CUDA graph mode, updates the static index + scatter buffers via + ``.copy_()`` so the captured replay reads stable addresses. + """ + self.dtype = dtype + self.attn_wrapper.plan( + qo_indptr, + kv_indptr, + kv_indices, + kv_len_arr, + self.num_heads, + self.head_dim_ckv, + self.head_dim_kpe, + self.page_size, + causal, + self.sm_scale, + dtype, + dtype, + ) + + # Per-token (page, offset) for the latent scatter, mirroring + # FlashInferPrefillWrapper.plan: token j of request r lands at absolute + # position g = old_len_r + j -> page kv_indices[kv_indptr[r] + g//ps], + # offset g % ps. old_len_r = kv_len_arr[r] - new_tokens_r. + n_req = qo_indptr.shape[0] - 1 + starts = qo_indptr[:-1].to(torch.int32) + lens = (qo_indptr[1:] - qo_indptr[:-1]).to(torch.int32) + total_tokens = int(lens.sum().item()) + self._total_tokens = total_tokens + + seg = torch.repeat_interleave( + torch.arange(n_req, dtype=torch.int32, device=self.device), lens + ) + intra = torch.arange( + total_tokens, dtype=torch.int32, device=self.device + ) - torch.repeat_interleave(starts, lens) + + start_new = kv_len_arr[seg] - lens[seg] + g = start_new + intra + + page_off = torch.div(g, self.page_size, rounding_mode="floor").to(torch.int32) + off_in_page = (g - page_off * self.page_size).to(torch.int32) + abs_page_ptr = kv_indptr[:-1][seg] + page_off + + token_to_page = kv_indices[abs_page_ptr].to(torch.long) + token_to_cache = off_in_page.to(torch.long) + + if self.use_cuda_graph: + self.token_to_page[:total_tokens].copy_(token_to_page) + self.token_to_cache[:total_tokens].copy_(token_to_cache) + if total_tokens < self.max_total_tokens: + self.token_to_page[total_tokens:] = 0 + self.token_to_cache[total_tokens:] = 0 + else: + self.token_to_page = token_to_page + self.token_to_cache = token_to_cache + + @torch.compiler.disable + def set_latent(self, latent_cache_layer: torch.Tensor, latent: torch.Tensor): + """Scatter one compressed latent per new token into the paged cache. + + Args: + latent_cache_layer: [max_pages, page_size, ckv + kpe] + latent: [total_tokens, ckv + kpe] = ``cat([kv_c, k_pe])`` for the new tokens. + """ + n = self._total_tokens + page_idx = self.token_to_page[:n] + cache_idx = self.token_to_cache[:n] + latent_cache_layer[page_idx, cache_idx] = latent[:n].to(latent_cache_layer.dtype) + + @torch.compiler.disable + def run( + self, + q_nope: torch.Tensor, + q_pe: torch.Tensor, + ckv_cache: torch.Tensor, + kpe_cache: torch.Tensor, + ) -> torch.Tensor: + """Run the planned MLA kernel. + + Args: + q_nope: [T, H, ckv] + q_pe: [T, H, kpe] + ckv_cache: [max_pages, page_size, ckv] (strided view of the latent cache) + kpe_cache: [max_pages, page_size, kpe] (strided view of the latent cache) + Returns: + [T, H, ckv] + """ + return self.attn_wrapper.run( + q_nope.to(self.dtype), q_pe.to(self.dtype), + ckv_cache, kpe_cache, return_lse=False, + ) diff --git a/test/integration/test_kimi_mla_absorb_forward.py b/test/integration/test_kimi_mla_absorb_forward.py new file mode 100644 index 000000000..45d3d5da0 --- /dev/null +++ b/test/integration/test_kimi_mla_absorb_forward.py @@ -0,0 +1,158 @@ +"""Phase-A GPU check: the REAL ``KimiMLAAttention.forward`` absorbed branch. + +The CPU gate (``test/modular/test_kimi_mla_absorb.py``) proves the absorption +algebra with pure-torch references. This test drives the actual wired forward — +``config.mla_absorb=True`` -> ``_forward_absorbed`` -> ``run_attention_mla`` — so +it needs a GPU (MLA RMSNorm uses a FlashInfer kernel). A ``_MockMLALatentCache`` +stands in for the Phase-B paged latent backend: its ``run_attention_mla`` does a +causal SDPA over ``[kv_c | k_pe]`` (value = ``kv_c``) at the DeepSeek scale, which +is exactly what the FlashInfer MLA kernel will compute. The real kernel is locked +to ckv=512/kpe=64 so it can't run at the reduced dims — that path is validated in +Phase B on real dims. + +Matching the independent DeepSeek MLA (materialized k_nope/v, no absorption) proves +the wired absorbed forward is numerically the naive path. + +Run: pytest test/integration/test_kimi_mla_absorb_forward.py -v +""" +import pytest +import torch +import torch.nn.functional as F + +from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention +from mstar.model.kimi_k2_7.components.rope import ( + _yarn_find_correction_range, + _yarn_linear_ramp_mask, + rotate_gptj, + yarn_get_mscale, +) +from mstar.model.kimi_k2_7.config import KimiK2Config + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="the absorbed forward runs MLA RMSNorm (a FlashInfer kernel) on GPU", +) + +DEVICE = "cuda" + + +# -------------------------------------------------------------------------- +# References (device-aware; mirror test_kimi_mla_paged.py). +# -------------------------------------------------------------------------- + +def _ref_rmsnorm(x, weight, eps): + x32 = x.float() + x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) + return weight * x32.to(x.dtype) + + +def _ref_yarn_rope(pos, q_pe, k_pe, cfg): + r = cfg.rope_scaling + rotary_dim, base, factor = cfg.qk_rope_head_dim, cfg.rope_theta, r["factor"] + max_pos = r["original_max_position_embeddings"] + beta_fast, beta_slow = r.get("beta_fast", 32), r.get("beta_slow", 1) + mscale, mscale_all_dim = r.get("mscale", 1.0), r.get("mscale_all_dim", 0.0) + pos_freqs = base ** (torch.arange(0, rotary_dim, 2, device=q_pe.device).float() / rotary_dim) + ext, interp = 1.0 / pos_freqs, 1.0 / (factor * pos_freqs) + low, high = _yarn_find_correction_range(beta_fast, beta_slow, rotary_dim, base, max_pos) + mask = 1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float).to(q_pe.device) + inv_freq = interp * (1 - mask) + ext * mask + amp = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim) + freqs = torch.outer(pos.float(), inv_freq) + cos = (freqs.cos() * amp).repeat_interleave(2, -1).unsqueeze(-2) + sin = (freqs.sin() * amp).repeat_interleave(2, -1).unsqueeze(-2) + qr = q_pe.float() * cos + rotate_gptj(q_pe.float()) * sin + kr = k_pe.float() * cos + rotate_gptj(k_pe.float()) * sin + return qr.to(q_pe.dtype), kr.to(k_pe.dtype) + + +def _sdpa_causal(q, k, v, scale): + qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) # (H,T,D) + t = q.shape[0] + causal = torch.triu(torch.full((t, t), float("-inf"), device=q.device), diagonal=1) + attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1) + return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) + + +def _deepseek_scale(cfg): + r = cfg.rope_scaling + mscale = yarn_get_mscale(r["factor"], r.get("mscale_all_dim", 0.0)) + return cfg.qk_head_dim ** -0.5 * mscale * mscale + + +def _ref_deepseek_mla(attn, cfg, h, pos): + """Naive DeepSeek MLA (materialized k_nope/v; no absorption) — ground truth.""" + t, heads = h.shape[0], attn.num_heads + d_nope, d_rope, d_v, latent = ( + cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank) + eps = cfg.rms_norm_eps + q = _ref_rmsnorm(F.linear(h, attn.q_a_proj.weight), attn.q_a_layernorm.weight, eps) + q = F.linear(q, attn.q_b_proj.weight).view(t, heads, cfg.qk_head_dim) + q_nope, q_pe = q.split([d_nope, d_rope], dim=-1) + lat = F.linear(h, attn.kv_a_proj_with_mqa.weight) + kv_a, k_pe = lat.split([latent, d_rope], dim=-1) + kv = F.linear(_ref_rmsnorm(kv_a, attn.kv_a_layernorm.weight, eps), + attn.kv_b_proj.weight).view(t, heads, d_nope + d_v) + k_nope, v = kv.split([d_nope, d_v], dim=-1) + k_pe = k_pe.view(t, 1, d_rope) + q_pe, k_pe = _ref_yarn_rope(pos, q_pe, k_pe, cfg) + q = torch.cat([q_nope, q_pe], dim=-1) + k = torch.cat([k_nope, k_pe.expand(t, heads, d_rope)], dim=-1) + out = _sdpa_causal(q, k, v, _deepseek_scale(cfg)).reshape(t, heads * d_v) + return F.linear(out, attn.o_proj.weight) + + +class _MockMLALatentCache: + """Phase-B latent-backend stand-in: causal SDPA over [kv_c|k_pe], value=kv_c.""" + + def __init__(self, sm_scale): + self.sm_scale = sm_scale + + def set_layer_idx(self, _i): + pass + + def set_active_label(self, _l): + pass + + def advance_seq_lens(self, *_a, **_k): + pass + + def run_attention_mla(self, q_nope, q_pe, kv_c, k_pe): + t, heads, latent = q_nope.shape + d_rope = q_pe.shape[-1] + query = torch.cat([q_nope, q_pe], dim=-1) # (T,H,L+Drope) + kv_c_h = kv_c.expand(t, heads, latent) # MQA broadcast + key = torch.cat([kv_c_h, k_pe.expand(t, heads, d_rope)], dim=-1) + return _sdpa_causal(query, key, kv_c_h, self.sm_scale) # (T,H,L) + + +def _build_attention(cfg, dtype): + attn = KimiMLAAttention(cfg).to(device=DEVICE, dtype=dtype) + for lin in (attn.q_a_proj, attn.q_b_proj, attn.kv_a_proj_with_mqa, + attn.kv_b_proj, attn.o_proj): + lin.weight.data.normal_(0, 0.03) + for norm in (attn.q_a_layernorm, attn.kv_a_layernorm): + norm.weight.data.normal_(1.0, 0.02) + attn.process_weights_after_loading() # build w_kc / w_vc on-device + return attn + + +def test_absorbed_forward_matches_deepseek(): + torch.manual_seed(0) + cfg = KimiK2Config.reduced() + cfg.mla_absorb = True + dtype = torch.bfloat16 + attn = _build_attention(cfg, dtype) + assert attn.w_kc is not None and attn.w_vc is not None + + t = 7 + h = torch.randn(t, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.1 + pos = torch.arange(t, device=DEVICE) + + cache = _MockMLALatentCache(_deepseek_scale(cfg)) + with torch.no_grad(): + got = attn(h, cache, pos) + + expected = _ref_deepseek_mla(attn, cfg, h, pos) + assert got.shape == (t, cfg.hidden_size) + torch.testing.assert_close(got, expected, rtol=3e-2, atol=3e-2) diff --git a/test/integration/test_kimi_mla_absorb_kernel.py b/test/integration/test_kimi_mla_absorb_kernel.py new file mode 100644 index 000000000..ee803cc6d --- /dev/null +++ b/test/integration/test_kimi_mla_absorb_kernel.py @@ -0,0 +1,352 @@ +"""Phase-B follow-up GPU check: the FlashInfer MLA **kernel** fast path. + +``MlaAbsorbCacheManager`` uses ``flashinfer.mla.BatchMLAPagedAttentionWrapper`` +(the dedicated ckv=512/kpe=64 kernel) instead of the SDPA gather loop whenever +``_mla_kernel_available`` says the dims + GPU support it (real Kimi dims on sm90). +This drives the manager at real dims so the kernel activates and asserts: + + kernel path == SDPA path == independent accumulate-everything reference + +across a multi-page prefill, a decode step, and a batched (multi-request) decode. +The kernel manager (``mla_ckv_dim=512``) and the SDPA manager (``mla_ckv_dim=None``, +which forces the fallback at the SAME dims) share identical random inputs, so any +divergence is the kernel's. It also asserts the probe declines reduced dims (the +kernel is dim-locked and crashes off-dim, so reduced configs MUST use SDPA). + +Requires a Hopper (sm90) GPU — off sm90 the probe declines and the "kernel" +manager silently uses SDPA, making the comparison a tautology, so we skip. + +Run: pytest test/integration/test_kimi_mla_absorb_kernel.py -v +""" +import pytest +import torch + +from mstar.communication.tensors import LocalTransferEngine +from mstar.engine.cache_manager import ( + MlaAbsorbCacheManager, + _mla_kernel_available, + create_cache_manager, +) +from mstar.engine.kv_store import ( + KVCacheConfig, + PagedAllocationManager, + TransferEngineInfo, +) + +_IS_SM90 = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 9 + +pytestmark = pytest.mark.skipif( + not _IS_SM90, + reason="the FlashInfer MLA kernel fast path requires a Hopper (sm90) GPU", +) + +DEVICE = torch.device("cuda") +# Real Kimi latent dims — the only dims the MLA kernel supports. +L, DROPE = 512, 64 + + +# -------------------------------------------------------------------------- +# Real paged latent cache manager. ``ckv_dim=512`` -> kernel; None -> SDPA. +# -------------------------------------------------------------------------- + +def _make_cm(softmax_scale, ckv_dim, request_ids, page_size=4, max_num_pages=64): + latent_width = L + DROPE + kv_cache = torch.zeros( + 2, max_num_pages, page_size, latent_width, dtype=torch.bfloat16, device=DEVICE, + ).contiguous() + kv_cfg = KVCacheConfig( + num_layers=2, num_kv_heads=1, head_dim=latent_width, + max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages, + page_size=page_size, num_qo_heads=1, + attention_backend="mla_absorb", softmax_scale=softmax_scale, + mla_ckv_dim=ckv_dim, + ) + transfer_info = TransferEngineInfo( + my_entity_id="kimi_mla_kernel_test", + my_session_id="kimi_session", + transfer_engine=LocalTransferEngine("localhost"), + ) + alloc = PagedAllocationManager( + config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info, + ) + for rid in request_ids: + alloc.add_request(rid, ["main"]) + from mstar.engine.cache_manager import WorkspaceBufferManager + buffers = WorkspaceBufferManager(128 * 1024 * 1024, device=DEVICE) + cm = create_cache_manager( + request_ids=list(request_ids), + active_labels_per_request={rid: "main" for rid in request_ids}, + kv_cache=kv_cache, + alloc_manager=alloc, + buffer_manager=buffers, + kv_cache_config=kv_cfg, + device=DEVICE, + ) + assert isinstance(cm, MlaAbsorbCacheManager) + cm.set_active_label("main") + cm.set_layer_idx(0) + return cm, alloc + + +# -------------------------------------------------------------------------- +# Independent reference: accumulate every (kv_c, k_pe), causal SDPA (no paging). +# -------------------------------------------------------------------------- + +def _ref_mla_step(q_nope_new, q_pe_new, kv_c_all, k_pe_all, scale): + """Intended MLA output for the NEW query tokens. ``kv_c_all``/``k_pe_all`` are + EVERY latent cached so far (including this step's); query j sits at absolute + position old_len+j and attends cached 0..old_len+j; value = the kv_c part.""" + sl, _H, _L = q_nope_new.shape + total = kv_c_all.shape[0] + old_len = total - sl + query = torch.cat([q_nope_new, q_pe_new], dim=-1) # [sl,H,L+Drope] + key = torch.cat([kv_c_all.squeeze(1), k_pe_all.squeeze(1)], dim=-1) # [total,L+Drope] + value = kv_c_all.squeeze(1) # [total,L] + qt = query.transpose(0, 1).float() # [H,sl,L+Drope] + scores = torch.einsum("hqd,kd->hqk", qt, key.float()) * scale + q_pos = old_len + torch.arange(sl, device=DEVICE) + k_pos = torch.arange(total, device=DEVICE) + mask = torch.where(k_pos[None, :] <= q_pos[:, None], 0.0, + torch.tensor(float("-inf"), device=DEVICE)) + attn = (scores + mask).softmax(-1) + out = torch.einsum("hqk,kd->hqd", attn, value.float()) # [H,sl,L] + return out.transpose(0, 1).to(q_nope_new.dtype) # [sl,H,L] + + +def _rand_step(sl, H, dtype=torch.bfloat16): + return ( + torch.randn(sl, H, L, device=DEVICE, dtype=dtype) * 0.1, + torch.randn(sl, H, DROPE, device=DEVICE, dtype=dtype) * 0.1, + torch.randn(sl, 1, L, device=DEVICE, dtype=dtype) * 0.1, + torch.randn(sl, 1, DROPE, device=DEVICE, dtype=dtype) * 0.1, + ) + + +def _run_single_req(cm, alloc, T, H, scale): + """Prefill T tokens (multi-page) + one decode step through ``cm``; return the + two outputs and the reference for each.""" + q_nope, q_pe, kv_c, k_pe = _rand_step(T, H) + cm.plan_attention(seq_lens=[T], is_causal=True, dtype=torch.bfloat16) + with torch.no_grad(): + got_prefill = cm.run_attention_mla(q_nope, q_pe, kv_c, k_pe) + torch.cuda.synchronize() + ref_prefill = _ref_mla_step(q_nope, q_pe, kv_c, k_pe, scale) + cm.advance_seq_lens() + + q_nope1, q_pe1, kv_c1, k_pe1 = _rand_step(1, H) + cm.plan_attention(seq_lens=[1], is_causal=True, dtype=torch.bfloat16) + with torch.no_grad(): + got_decode = cm.run_attention_mla(q_nope1, q_pe1, kv_c1, k_pe1) + torch.cuda.synchronize() + ref_decode = _ref_mla_step( + q_nope1, q_pe1, + torch.cat([kv_c, kv_c1], 0), torch.cat([k_pe, k_pe1], 0), scale, + ) + return (got_prefill, ref_prefill), (got_decode, ref_decode) + + +def test_kernel_matches_sdpa_and_reference(): + """Real dims (L=512, Drope=64, H=2): kernel == SDPA == reference, prefill+decode. + + Same seed/inputs into a kernel manager (``mla_ckv_dim=512``) and an SDPA + manager (``mla_ckv_dim=None``) at identical dims — both must match the + independent reference and each other.""" + H, T, ps = 2, 6, 4 # T=6 over page_size-4 = 2 pages + scale = (L + DROPE) ** -0.5 * 1.3 + assert _mla_kernel_available(L, DROPE, 9) is True + + torch.manual_seed(0) + cm_k, alloc_k = _make_cm(scale, ckv_dim=L, request_ids=["r0"], page_size=ps) + try: + (kp, refp), (kd, refd) = _run_single_req(cm_k, alloc_k, T, H, scale) + finally: + alloc_k.cleanup() + + torch.manual_seed(0) # identical inputs + cm_s, alloc_s = _make_cm(scale, ckv_dim=None, request_ids=["r0"], page_size=ps) + try: + (sp, _refp), (sd, _refd) = _run_single_req(cm_s, alloc_s, T, H, scale) + finally: + alloc_s.cleanup() + + # Sanity: the kernel manager actually took the kernel path (ps.wrapper set), + # the SDPA manager did not. + assert cm_k._plan_states["main"].wrapper is not None + assert cm_s._plan_states["main"].wrapper is None + + assert kp.shape == (T, H, L) and kd.shape == (1, H, L) + # kernel == reference + torch.testing.assert_close(kp, refp, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(kd, refd, rtol=2e-2, atol=2e-2) + # SDPA == reference + torch.testing.assert_close(sp, refp, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(sd, refd, rtol=2e-2, atol=2e-2) + # kernel == SDPA + torch.testing.assert_close(kp, sp, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(kd, sd, rtol=2e-2, atol=2e-2) + + +def test_kernel_batched_decode(): + """Batched (multi-request, varying lengths) decode through the kernel path. + + Two requests are prefilled to different lengths, then a single decode step + runs both together; each request's decode output must match its own + accumulate-everything reference.""" + H, ps = 2, 4 + scale = (L + DROPE) ** -0.5 + lens = [5, 9] # req0 prefill 5 tokens, req1 prefill 9 (spans >1 page) + rids = ["r0", "r1"] + + torch.manual_seed(1) + cm, alloc = _make_cm(scale, ckv_dim=L, request_ids=rids, page_size=ps) + try: + # ---- prefill both requests (concatenated queries) ---- + pref = [_rand_step(sl, H) for sl in lens] + q_nope = torch.cat([p[0] for p in pref], 0) + q_pe = torch.cat([p[1] for p in pref], 0) + kv_c = torch.cat([p[2] for p in pref], 0) + k_pe = torch.cat([p[3] for p in pref], 0) + cm.plan_attention(seq_lens=lens, is_causal=True, dtype=torch.bfloat16) + with torch.no_grad(): + cm.run_attention_mla(q_nope, q_pe, kv_c, k_pe) + torch.cuda.synchronize() + cm.advance_seq_lens() + assert cm._plan_states["main"].wrapper is not None + + # ---- one decode step for both requests ---- + dec = [_rand_step(1, H) for _ in lens] + dq_nope = torch.cat([d[0] for d in dec], 0) + dq_pe = torch.cat([d[1] for d in dec], 0) + dkv_c = torch.cat([d[2] for d in dec], 0) + dk_pe = torch.cat([d[3] for d in dec], 0) + cm.plan_attention(seq_lens=[1, 1], is_causal=True, dtype=torch.bfloat16) + with torch.no_grad(): + got = cm.run_attention_mla(dq_nope, dq_pe, dkv_c, dk_pe) + torch.cuda.synchronize() + assert got.shape == (2, H, L) + + # ---- per-request reference (its prefill latents + its decode latent) ---- + for i in range(2): + kv_c_all = torch.cat([pref[i][2], dec[i][2]], 0) # [len_i+1,1,L] + k_pe_all = torch.cat([pref[i][3], dec[i][3]], 0) + ref = _ref_mla_step(dec[i][0], dec[i][1], kv_c_all, k_pe_all, scale) + torch.testing.assert_close(got[i:i + 1], ref, rtol=2e-2, atol=2e-2) + finally: + alloc.cleanup() + + +def test_mla_wrapper_cuda_graph_capture_replay(): + """The MLA kernel wrapper + latent scatter are CUDA-graph capturable. + + Captures ONE decode graph (``set_latent`` + ``run`` over static buffers) then + replays it across two decode steps — re-planning each step (which updates the + wrapper's static index/scatter buffers via ``.copy_()``, advancing kv_len + + scatter offset) and copying fresh queries/latents into the static input + buffers before replay. Each replay's output must match the independent + accumulate-everything reference for that step. This mirrors exactly what + ``CudaGraphRunner`` does across decode steps, proving the absorbed-decode + primitive captures + replays correctly with a moving page table. + """ + from mstar.utils.flashinfer_utils import FlashInferMLAWrapper + + bs, H, ps, max_pages = 2, 2, 4, 64 + latent_w = L + DROPE + scale = latent_w ** -0.5 * 1.1 + dtype = torch.bfloat16 + + cache = torch.zeros(max_pages, ps, latent_w, device=DEVICE, dtype=dtype) + ws = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=DEVICE) + wrapper = FlashInferMLAWrapper( + ws, num_heads=H, head_dim_ckv=L, head_dim_kpe=DROPE, page_size=ps, + sm_scale=scale, batch_size=bs, max_num_pages=max_pages, + max_total_tokens=bs, device=DEVICE, use_cuda_graph=True, + ) + + # Per-request fixed page ranges + prefill lengths. prefill=6/9 with ps=4 keep + # the two following decode steps within the same 2/3 pages (kv_indices stable; + # only kv_len + the scatter offset advance) — the common decode case. + req_pages = [[0, 1], [2, 3, 4]] + prefill = [6, 9] + torch.manual_seed(7) + prefix = [] # [prefill_r, latent_w] latents already resident per request + for r in range(bs): + lat = torch.randn(prefill[r], latent_w, device=DEVICE, dtype=dtype) * 0.1 + for t in range(prefill[r]): + cache[req_pages[r][t // ps], t % ps] = lat[t] + prefix.append(lat) + + kv_indptr = torch.tensor( + [0, len(req_pages[0]), len(req_pages[0]) + len(req_pages[1])], + device=DEVICE, dtype=torch.int32, + ) + kv_indices = torch.tensor(req_pages[0] + req_pages[1], device=DEVICE, dtype=torch.int32) + qo_indptr = torch.tensor([0, 1, 2], device=DEVICE, dtype=torch.int32) + + # Static input buffers replay reads from. + q_nope_s = torch.zeros(bs, H, L, device=DEVICE, dtype=dtype) + q_pe_s = torch.zeros(bs, H, DROPE, device=DEVICE, dtype=dtype) + latent_s = torch.zeros(bs, latent_w, device=DEVICE, dtype=dtype) + + def plan_step(step): # step 1 -> position prefill_r, step 2 -> prefill_r+1 + kv_len_arr = torch.tensor( + [prefill[r] + step for r in range(bs)], device=DEVICE, dtype=torch.int32, + ) + wrapper.plan(qo_indptr, kv_indptr, kv_indices, kv_len_arr, causal=True, dtype=dtype) + + def fill_inputs(seed): + torch.manual_seed(seed) + q_nope_s.copy_(torch.randn(bs, H, L, device=DEVICE, dtype=dtype) * 0.1) + q_pe_s.copy_(torch.randn(bs, H, DROPE, device=DEVICE, dtype=dtype) * 0.1) + latent_s.copy_(torch.randn(bs, latent_w, device=DEVICE, dtype=dtype) * 0.1) + + # ---- warmup (side stream) then capture ONE decode graph ---- + plan_step(1) + fill_inputs(100) + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + wrapper.set_latent(cache, latent_s) + wrapper.run(q_nope_s, q_pe_s, cache[..., :L], cache[..., L:]) + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + wrapper.set_latent(cache, latent_s) + out_static = wrapper.run(q_nope_s, q_pe_s, cache[..., :L], cache[..., L:]) + + # ---- replay two decode steps; each must match its accumulate reference ---- + decode_hist = [[] for _ in range(bs)] # decode latents scattered so far, per req + prev = None + for step in (1, 2): + fill_inputs(step) # fresh queries + decode latents into static bufs + plan_step(step) # re-plan: advance kv_len + scatter offset + g.replay() + torch.cuda.synchronize() + for r in range(bs): + decode_hist[r].append(latent_s[r:r + 1].clone()) + got = out_static.clone() + + # reference: query attends [prefix_r ; all decode latents so far] + ref = torch.empty(bs, H, L, device=DEVICE, dtype=dtype) + for r in range(bs): + all_lat = torch.cat([prefix[r]] + decode_hist[r], dim=0) # [prefill+step, w] + ref[r:r + 1] = _ref_mla_step( + q_nope_s[r:r + 1], q_pe_s[r:r + 1], + all_lat[:, :L].unsqueeze(1), all_lat[:, L:].unsqueeze(1), scale, + ) + torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2) + if prev is not None: + # Sanity: distinct inputs -> distinct outputs (values really flow + # through the static buffers, not baked into the captured graph). + assert not torch.allclose(got, prev, atol=1e-3) + prev = got + + +def test_probe_declines_reduced_dims(): + """The kernel is dim-locked (ckv=512/kpe=64) and crashes off-dim, so the probe + must decline reduced dims / pre-sm90 (→ SDPA), and accept only real dims on sm90.""" + assert _mla_kernel_available(L, DROPE, 9) is True # real dims, Hopper + assert _mla_kernel_available(32, 8, 9) is False # reduced dims + assert _mla_kernel_available(L, DROPE, 8) is False # pre-sm90 + assert _mla_kernel_available(L, DROPE, 10) is False # Blackwell (wants trtllm path) diff --git a/test/integration/test_kimi_mla_absorb_marlin_merge.py b/test/integration/test_kimi_mla_absorb_marlin_merge.py new file mode 100644 index 000000000..c699c91a3 --- /dev/null +++ b/test/integration/test_kimi_mla_absorb_marlin_merge.py @@ -0,0 +1,119 @@ +"""Merge check: the absorbed-MLA and Marlin-MoE post-load hooks COMPOSE. + +The MLA weight-absorption work and the Marlin W4A16 MoE work were developed in +separate silos; both hang a ``process_weights_after_loading`` finalizer off the +same generic walker (``mstar.model.components.quantization.process_weights_after_loading``, +called once in ``kimi_model.py::_create_submodule``). This test builds a tree +holding BOTH a weight-absorbed ``KimiMLAAttention`` and a Marlin-backed +``KimiSparseMoeBlock`` and runs that single walker pass, asserting it finalizes +BOTH: the attention's ``w_kc``/``w_vc``/``fused_qkv_a_proj`` and the MoE's Marlin +repack (+ freed packed params). This is the production-default combination +(``k27_code`` => ``mla_absorb=True`` + ``quant_kernel="auto"`` => Marlin on sm80+), +which neither silo validated together. A Marlin-MoE forward post-walker confirms +the block still runs in the combined module. (Reduced dims => the attention uses +the SDPA-over-latent path, not the dim-locked FlashInfer MLA kernel; the kernel + +Marlin combined forward at real dims is covered by the TP8 serve smoke.) + +Run: pytest test/integration/test_kimi_mla_absorb_marlin_merge.py -v +""" +import pytest +import torch +from torch import nn + +from mstar.model.components.quantization import process_weights_after_loading +from mstar.model.kimi_k2_7._testing import fake_quantize_weight +from mstar.model.kimi_k2_7.config import KimiK2Config + +pytestmark = pytest.mark.skipif( + not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8), + reason="the Marlin backend needs a CUDA GPU with sm80+", +) + +DEVICE = "cuda" +GROUP_SIZE = 32 +PACK_FACTOR = 8 + + +def _quantize_stack(weight): + """Per-expert fake-quantize a bf16 stack to packed int32 + bf16 scales.""" + E, N, K = weight.shape + packed = torch.empty((E, N, K // PACK_FACTOR), dtype=torch.int32, device=DEVICE) + scale = torch.empty((E, N, K // GROUP_SIZE), dtype=torch.bfloat16, device=DEVICE) + for e in range(E): + p, s, _d = fake_quantize_weight( + weight[e], num_bits=4, group_size=GROUP_SIZE, symmetric=True, + scale_dtype=torch.bfloat16, + ) + packed[e], scale[e] = p.to(DEVICE), s.to(DEVICE) + return packed, scale + + +def _build_marlin_moe(cfg): + """A tp=1 Marlin-legal KimiSparseMoeBlock with synthetic packed experts + (mirrors test_kimi_moe_marlin._build_block).""" + from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock + + with torch.device("meta"): + block = KimiSparseMoeBlock(cfg) + block = block.to(torch.bfloat16) + block.to_empty(device=DEVICE) + for p in block.parameters(): + if p.dtype.is_floating_point: + with torch.no_grad(): + p.copy_(torch.randn_like(p) * 0.1) + E, H, I = cfg.n_routed_experts, cfg.hidden_size, cfg.moe_intermediate_size + w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16) + w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16) + w1_packed, w1_scale = _quantize_stack(w1) + w2_packed, w2_scale = _quantize_stack(w2) + with torch.no_grad(): + block.experts.gate_up_proj_packed.data = w1_packed + block.experts.gate_up_proj_scale.data = w1_scale + block.experts.down_proj_packed.data = w2_packed + block.experts.down_proj_scale.data = w2_scale + return block + + +def _build_absorbed_attn(cfg): + """A materialized weight-absorbed KimiMLAAttention (random weights on CUDA).""" + from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention + + torch.manual_seed(0) + return KimiMLAAttention(cfg).to(device=DEVICE, dtype=torch.bfloat16) + + +def test_generic_walker_composes_absorbed_mla_and_marlin_moe(): + """One ``process_weights_after_loading`` pass finalizes BOTH the absorbed MLA + projections and the Marlin MoE repack — the merged production default.""" + cfg = KimiK2Config.reduced_marlin() + cfg.mla_absorb = True # absorbed attention + Marlin experts in one config + + attn = _build_absorbed_attn(cfg) + moe = _build_marlin_moe(cfg) + root = nn.Module() + root.add_module("attn", attn) + root.add_module("moe", moe) + + # --- pre-walker: neither finalizer has run --- + assert attn.mla_absorb + assert attn.w_kc is None and attn.w_vc is None and attn.fused_qkv_a_proj_weight is None + assert moe.experts.gate_up_proj_packed.numel() > 0 + assert not getattr(moe, "_use_marlin", False) + + # --- one generic walker pass (the exact call kimi_model._create_submodule makes) --- + process_weights_after_loading(root, torch.device(DEVICE)) + + # --- post-walker: absorbed-MLA finalizer ran --- + assert attn.w_kc is not None, "walker did not build the absorbed w_kc" + assert attn.w_vc is not None, "walker did not build the absorbed w_vc" + assert attn.fused_qkv_a_proj_weight is not None, "walker did not build fused_qkv_a_proj" + + # --- post-walker: Marlin-MoE finalizer ran --- + assert moe._use_marlin, "walker did not select the Marlin backend" + assert moe.experts.gate_up_proj_packed.numel() == 0, "packed experts not freed after repack" + + # --- the Marlin MoE forward still runs in the combined module --- + x = torch.randn(4, cfg.hidden_size, device=DEVICE, dtype=torch.bfloat16) * 0.1 + with torch.no_grad(): + out = moe(x) + assert out.shape == x.shape and torch.isfinite(out).all() diff --git a/test/integration/test_kimi_mla_absorb_paged.py b/test/integration/test_kimi_mla_absorb_paged.py new file mode 100644 index 000000000..6e83307e3 --- /dev/null +++ b/test/integration/test_kimi_mla_absorb_paged.py @@ -0,0 +1,182 @@ +"""Phase-B GPU check: the REAL paged compressed-latent MLA backend. + +Drives ``MlaAbsorbCacheManager.run_attention_mla`` over a genuine 4D latent +paged cache (real ``PagedAllocationManager`` + ``create_cache_manager``), with +SYNTHETIC random ``q_nope/q_pe/kv_c/k_pe`` — no ``KimiMLAAttention`` needed. The +manager writes ``cat([kv_c, k_pe])`` as one latent vector per token into the +paged cache at its (page, offset), then per request gathers its full cached +latent and runs a causal SDPA (query = ``cat([q_nope, q_pe])``, key = the full +latent, value = its first ``L`` dims) at ``kv_cache_config.softmax_scale``. + +The reference is INDEPENDENT of the paging machinery: it accumulates every +``(kv_c, k_pe)`` seen so far into contiguous tensors and runs the same causal +SDPA (no pages, no scatter/gather). Matching it across a multi-page prefill and a +following decode step proves the paged scatter/gather + causal mask + scale are +correct — at real Kimi latent dims (L=512, Drope=64) and reduced dims. + +Run: pytest test/integration/test_kimi_mla_absorb_paged.py -v +""" +import pytest +import torch + +from mstar.communication.tensors import LocalTransferEngine +from mstar.engine.cache_manager import ( + MlaAbsorbCacheManager, + WorkspaceBufferManager, + create_cache_manager, +) +from mstar.engine.kv_store import ( + KVCacheConfig, + PagedAllocationManager, + TransferEngineInfo, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="the paged compressed-latent MLA backend runs on GPU", +) + +DEVICE = torch.device("cuda") + + +# -------------------------------------------------------------------------- +# Real paged latent cache manager (mirrors test_kimi_mla_paged.py, but 4D). +# -------------------------------------------------------------------------- + +def _make_latent_cache_manager( + latent_width, dtype, softmax_scale, page_size=4, max_num_pages=64 +): + # 4D latent cache: [num_layers, max_pages, page_size, latent_width] + # (the shape KVCacheEngine.load_model allocates for attention_backend + # "mla_absorb"). Small page_size so a handful of tokens spans >1 page. + kv_cache = torch.zeros( + 2, max_num_pages, page_size, latent_width, + dtype=dtype, device=DEVICE, + ).contiguous() + kv_cfg = KVCacheConfig( + num_layers=2, num_kv_heads=1, head_dim=latent_width, + max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages, + page_size=page_size, num_qo_heads=1, + attention_backend="mla_absorb", softmax_scale=softmax_scale, + ) + transfer_info = TransferEngineInfo( + my_entity_id="kimi_mla_absorb_test", + my_session_id="kimi_session", + transfer_engine=LocalTransferEngine("localhost"), + ) + alloc = PagedAllocationManager( + config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info, + ) + alloc.add_request("r0", ["main"]) + buffers = WorkspaceBufferManager(64 * 1024 * 1024, device=DEVICE) + cm = create_cache_manager( + request_ids=["r0"], + active_labels_per_request={"r0": "main"}, + kv_cache=kv_cache, + alloc_manager=alloc, + buffer_manager=buffers, + kv_cache_config=kv_cfg, + device=DEVICE, + ) + assert isinstance(cm, MlaAbsorbCacheManager) + return cm, alloc + + +# -------------------------------------------------------------------------- +# Independent reference: accumulate all (kv_c, k_pe), causal SDPA (no paging). +# -------------------------------------------------------------------------- + +def _ref_mla_step(q_nope_new, q_pe_new, kv_c_all, k_pe_all, scale): + """The intended MLA output for the NEW query tokens. + + q_nope_new [sl,H,L], q_pe_new [sl,H,Drope] are just this step's queries; + kv_c_all [total,1,L], k_pe_all [total,1,Drope] are EVERY latent cached so far + (including this step's). Query j sits at absolute position old_len+j and + attends to cached 0..old_len+j; value is the kv_c (first L) part. + """ + sl, _H, L = q_nope_new.shape + total = kv_c_all.shape[0] + old_len = total - sl + + query = torch.cat([q_nope_new, q_pe_new], dim=-1) # [sl,H,L+Drope] + kv_c_h = kv_c_all.squeeze(1) # [total,L] + k_pe_h = k_pe_all.squeeze(1) # [total,Drope] + key = torch.cat([kv_c_h, k_pe_h], dim=-1) # [total,L+Drope] + value = kv_c_h # [total,L] + + qt = query.transpose(0, 1).float() # [H,sl,L+Drope] + scores = torch.einsum("hqd,kd->hqk", qt, key.float()) * scale # [H,sl,total] + q_pos = old_len + torch.arange(sl, device=DEVICE) + k_pos = torch.arange(total, device=DEVICE) + mask = torch.where( + k_pos[None, :] <= q_pos[:, None], + 0.0, + torch.tensor(float("-inf"), device=DEVICE), + ) + attn = (scores + mask).softmax(-1) + out = torch.einsum("hqk,kd->hqd", attn, value.float()) # [H,sl,L] + return out.transpose(0, 1).to(q_nope_new.dtype) # [sl,H,L] + + +def _rand_step(sl, H, L, Drope, dtype): + return ( + torch.randn(sl, H, L, device=DEVICE, dtype=dtype) * 0.1, + torch.randn(sl, H, Drope, device=DEVICE, dtype=dtype) * 0.1, + torch.randn(sl, 1, L, device=DEVICE, dtype=dtype) * 0.1, + torch.randn(sl, 1, Drope, device=DEVICE, dtype=dtype) * 0.1, + ) + + +def _run_prefill_then_decode(L, Drope, H, T, page_size): + """Drive a multi-page prefill + a decode step through the real paged backend + and compare each to the independent accumulate-everything reference.""" + torch.manual_seed(0) + dtype = torch.bfloat16 + # Arbitrary MLA-style scale (the backend must apply exactly this value). + scale = (L + Drope) ** -0.5 * 1.3 + + cm, alloc = _make_latent_cache_manager(L + Drope, dtype, scale, page_size=page_size) + try: + cm.set_active_label("main") + cm.set_layer_idx(0) + + # ---- prefill: T tokens spanning more than one page ---- + assert T > page_size, "T must span >1 page to exercise page boundaries" + q_nope, q_pe, kv_c, k_pe = _rand_step(T, H, L, Drope, dtype) + cm.plan_attention(seq_lens=[T], is_causal=True, dtype=dtype) + with torch.no_grad(): + got_prefill = cm.run_attention_mla(q_nope, q_pe, kv_c, k_pe) + torch.cuda.synchronize() + + ref_prefill = _ref_mla_step(q_nope, q_pe, kv_c, k_pe, scale) + assert got_prefill.shape == (T, H, L) + torch.testing.assert_close(got_prefill, ref_prefill, rtol=2e-2, atol=2e-2) + + # Advance seq_len so the decode step sees the T cached tokens. + cm.advance_seq_lens() + + # ---- decode: 1 new token attends over all T+1 cached ---- + q_nope1, q_pe1, kv_c1, k_pe1 = _rand_step(1, H, L, Drope, dtype) + cm.plan_attention(seq_lens=[1], is_causal=True, dtype=dtype) + with torch.no_grad(): + got_decode = cm.run_attention_mla(q_nope1, q_pe1, kv_c1, k_pe1) + torch.cuda.synchronize() + + kv_c_all = torch.cat([kv_c, kv_c1], dim=0) # [T+1,1,L] + k_pe_all = torch.cat([k_pe, k_pe1], dim=0) # [T+1,1,Drope] + ref_decode = _ref_mla_step(q_nope1, q_pe1, kv_c_all, k_pe_all, scale) + assert got_decode.shape == (1, H, L) + torch.testing.assert_close(got_decode, ref_decode, rtol=2e-2, atol=2e-2) + finally: + alloc.cleanup() + + +def test_paged_latent_mla_real_dims(): + """Real Kimi MLA latent dims: L=512, Drope=64, H=2. Prefill (6 tokens over a + page_size-4 cache = 2 pages) + a decode step.""" + _run_prefill_then_decode(L=512, Drope=64, H=2, T=6, page_size=4) + + +def test_paged_latent_mla_reduced_dims(): + """Reduced dims: L=32, Drope=8, H=4. Same multi-page prefill + decode step.""" + _run_prefill_then_decode(L=32, Drope=8, H=4, T=6, page_size=4) diff --git a/test/integration/test_kimi_mla_absorb_serve.py b/test/integration/test_kimi_mla_absorb_serve.py new file mode 100644 index 000000000..b4c2da9ac --- /dev/null +++ b/test/integration/test_kimi_mla_absorb_serve.py @@ -0,0 +1,262 @@ +"""Absorbed-MLA serve smoke: the FULL model through the real mla_absorb backend. + +This is the end-to-end gate for the (now-default) weight-absorbed path. The +absorbed *math* is validated by ``test_kimi_mla_absorb.py`` (CPU) and the wired +attention forward by ``test_kimi_mla_absorb_forward.py``; the SDPA-over-latent +backend by ``test_kimi_mla_absorb_paged.py``. This test closes the loop: a whole +reduced ``KimiForCausalLM`` (all layers) built via the real +``get_submodule`` (meta -> to_empty -> load_weights -> post-load walker that builds +w_kc/w_vc + fused_qkv_a_proj) is driven for a prefill + several decode steps over a +genuine ``MlaAbsorbCacheManager`` + engine-shaped 4D latent cache. + +Correctness tie: the SAME synthetic checkpoint is loaded into a NAIVE reference +model (``mla_absorb=False``) and run through the mock cache at the DeepSeek scale; +the absorbed serve's prefill logits must match it (weight absorption is numerically +identical to naive up to fp rounding). This pins the absorbed serve path to the +M4-golden naive reference at full-model scale. + +Run: pytest test/integration/test_kimi_mla_absorb_serve.py -v +""" +import pytest +import torch + +from mstar.communication.tensors import LocalTransferEngine +from mstar.engine.cache_manager import WorkspaceBufferManager, create_cache_manager +from mstar.engine.kv_store import ( + KVCacheConfig, + PagedAllocationManager, + TransferEngineInfo, +) +from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM +from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock +from mstar.model.kimi_k2_7.components.rope import yarn_get_mscale +from mstar.model.kimi_k2_7.config import KimiK2Config +from mstar.model.kimi_k2_7.kimi_model import KimiK2Model +from mstar.model.kimi_k2_7.submodules import KimiLLMSubmodule +from mstar.model.submodule_base import ModelInputsFromEngine + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="absorbed serve smoke needs a GPU (real paged latent cache)", +) + +DEVICE = torch.device("cuda") + + +# -------------------------------------------------------------------------- +# Synthetic checkpoint (same serialization as test_kimi_submodule.py). +# -------------------------------------------------------------------------- + +def _fill_layer(layer, cfg): + a = layer.self_attn + for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): + lin.weight.data.normal_(0, 0.03) + for norm in (a.q_a_layernorm, a.kv_a_layernorm): + norm.weight.data.normal_(1.0, 0.02) + layer.input_layernorm.weight.data.normal_(1.0, 0.02) + layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02) + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + mlp.gate.weight.data.normal_(0, 1) + mlp.gate.e_score_correction_bias.data = torch.randn( + cfg.n_routed_experts, device=DEVICE, dtype=torch.float32) + mlp.experts.gate_up_proj.data.normal_(0, 0.05) + mlp.experts.down_proj.data.normal_(0, 0.05) + mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05) + else: + mlp.gate_up_proj.weight.data.normal_(0, 0.05) + mlp.down_proj.weight.data.normal_(0, 0.05) + + +def _build_reference(cfg): + """A naive-path (mla_absorb=False) model with random weights = the ground truth.""" + model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=torch.bfloat16) + model.model.embed_tokens.weight.data.normal_(0, 0.05) + model.model.norm.weight.data.normal_(1.0, 0.02) + model.lm_head.weight.data.normal_(0, 0.02) + for layer in model.model.layers: + _fill_layer(layer, cfg) + return model.eval() + + +def _hf_checkpoint(model, cfg): + inter = cfg.intermediate_size + moe_inter = cfg.moe_intermediate_size + shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts + m = model.model + sd = {"model.embed_tokens.weight": m.embed_tokens.weight} + for i, layer in enumerate(m.layers): + p = f"model.layers.{i}." + a = layer.self_attn + sd[p + "self_attn.q_a_proj.weight"] = a.q_a_proj.weight + sd[p + "self_attn.q_a_layernorm.weight"] = a.q_a_layernorm.weight + sd[p + "self_attn.q_b_proj.weight"] = a.q_b_proj.weight + sd[p + "self_attn.kv_a_proj_with_mqa.weight"] = a.kv_a_proj_with_mqa.weight + sd[p + "self_attn.kv_a_layernorm.weight"] = a.kv_a_layernorm.weight + sd[p + "self_attn.kv_b_proj.weight"] = a.kv_b_proj.weight + sd[p + "self_attn.o_proj.weight"] = a.o_proj.weight + sd[p + "input_layernorm.weight"] = layer.input_layernorm.weight + sd[p + "post_attention_layernorm.weight"] = layer.post_attention_layernorm.weight + mlp = layer.mlp + if isinstance(mlp, KimiSparseMoeBlock): + sd[p + "mlp.gate.weight"] = mlp.gate.weight + sd[p + "mlp.gate.e_score_correction_bias"] = mlp.gate.e_score_correction_bias + gup, dwn = mlp.experts.gate_up_proj, mlp.experts.down_proj + for e in range(cfg.n_routed_experts): + sd[p + f"mlp.experts.{e}.gate_proj.weight"] = gup[e, :moe_inter, :] + sd[p + f"mlp.experts.{e}.up_proj.weight"] = gup[e, moe_inter:, :] + sd[p + f"mlp.experts.{e}.down_proj.weight"] = dwn[e] + sh = mlp.shared_expert + sd[p + "mlp.shared_experts.gate_proj.weight"] = sh.gate_up_proj.weight[:shared_inter] + sd[p + "mlp.shared_experts.up_proj.weight"] = sh.gate_up_proj.weight[shared_inter:] + sd[p + "mlp.shared_experts.down_proj.weight"] = sh.down_proj.weight + else: + sd[p + "mlp.gate_proj.weight"] = mlp.gate_up_proj.weight[:inter] + sd[p + "mlp.up_proj.weight"] = mlp.gate_up_proj.weight[inter:] + sd[p + "mlp.down_proj.weight"] = mlp.down_proj.weight + sd["model.norm.weight"] = m.norm.weight + sd["lm_head.weight"] = model.lm_head.weight + return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()} + + +# -------------------------------------------------------------------------- +# Naive mock cache (DeepSeek-correct scale) for the reference logits. +# -------------------------------------------------------------------------- + +def _sdpa_causal(q, k, v, scale): + qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) + t = q.shape[0] + causal = torch.triu(torch.full((t, t), float("-inf"), device=q.device), diagonal=1) + attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1) + return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) + + +class _MockMLACache: + def __init__(self, head_dim): + self.scale = head_dim ** -0.5 + + def set_layer_idx(self, _i): + pass + + def advance_seq_lens(self, *_a, **_k): + pass + + def run_attention(self, q, k, v): + return _sdpa_causal(q, k, v, self.scale) + + +# -------------------------------------------------------------------------- +# Real mla_absorb backend (4D latent cache), engine-shaped. +# -------------------------------------------------------------------------- + +def _absorbed_softmax_scale(cfg): + r = cfg.rope_scaling + mscale = yarn_get_mscale(r["factor"], r.get("mscale_all_dim", 0.0)) + return cfg.qk_head_dim ** -0.5 * mscale * mscale + + +def _make_latent_cache_manager(cfg, dtype, page_size=128, max_num_pages=8): + latent_dim = cfg.kv_lora_rank + cfg.qk_rope_head_dim + kv_cache = torch.zeros( + cfg.num_hidden_layers, max_num_pages, page_size, latent_dim, + dtype=dtype, device=DEVICE, + ).contiguous() # 4D latent cache (matches KVCacheEngine's mla_absorb branch) + kv_cfg = KVCacheConfig( + num_layers=cfg.num_hidden_layers, num_kv_heads=1, head_dim=latent_dim, + max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages, + page_size=page_size, num_qo_heads=cfg.num_attention_heads, + attention_backend="mla_absorb", softmax_scale=_absorbed_softmax_scale(cfg), + ) + transfer_info = TransferEngineInfo( + my_entity_id="kimi_absorb_serve", my_session_id="kimi_session", + transfer_engine=LocalTransferEngine("localhost"), + ) + alloc = PagedAllocationManager( + config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info) + alloc.add_request("r0", ["main"]) + buffers = WorkspaceBufferManager(64 * 1024 * 1024, device=DEVICE) + cm = create_cache_manager( + request_ids=["r0"], active_labels_per_request={"r0": "main"}, + kv_cache=kv_cache, alloc_manager=alloc, buffer_manager=buffers, + kv_cache_config=kv_cfg, device=DEVICE, + ) + return cm, alloc + + +def _make_model(cfg, checkpoint_dir) -> KimiK2Model: + model = object.__new__(KimiK2Model) # skip __init__ (tokenizer/full config) + model.config = cfg + model.model_path_hf = str(checkpoint_dir) + model.cache_dir = None + model._submodule_cache = {} + return model + + +def _step(submodule, cm, graph_walk, token_ids): + engine_inputs = ModelInputsFromEngine( + request_ids=["r0"], per_request_info={}, cache_manager=cm) + ar_in = submodule.prepare_inputs( + graph_walk=graph_walk, fwd_info=None, inputs={"text_inputs": [token_ids]}) + packed = submodule.preprocess(graph_walk, engine_inputs, [ar_in]) + with torch.no_grad(): + out = submodule.forward(graph_walk, engine_inputs, **packed) + return out["logits"][0] + + +# -------------------------------------------------------------------------- +# Test +# -------------------------------------------------------------------------- + +def test_absorbed_serve_matches_naive_reference(tmp_path): + from safetensors.torch import save_file + + torch.manual_seed(0) + # One synthetic checkpoint, built from a naive reference model. + cfg_naive = KimiK2Config.reduced() # mla_absorb=False (reduced default) + assert cfg_naive.mla_absorb is False + ref = _build_reference(cfg_naive) + save_file(_hf_checkpoint(ref, cfg_naive), str(tmp_path / "model.safetensors")) + + # Reference logits: the naive model through the mock cache at the DeepSeek scale. + T = 6 + prompt = torch.randint(0, cfg_naive.vocab_size, (T,), device=DEVICE) + pos = torch.arange(T, device=DEVICE) + with torch.no_grad(): + ref_hidden = ref.model(prompt, _MockMLACache(cfg_naive.padded_head_dim), pos) + ref_logits = ref.lm_head(ref_hidden[-1:]) # (1, vocab) + + # Absorbed serve: load the SAME checkpoint into an mla_absorb model via the real + # build path; the post-load walker builds w_kc/w_vc + fused_qkv_a_proj. + cfg_absorb = KimiK2Config.reduced() + cfg_absorb.mla_absorb = True + model = _make_model(cfg_absorb, tmp_path) + submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16) + assert isinstance(submodule, KimiLLMSubmodule) + # Absorbed models DO carry derived buffers (w_kc/w_vc/fused, persistent=False). + buf_names = {n for n, _ in submodule.language_model.named_buffers()} + assert any("w_kc" in n for n in buf_names) and any("fused_qkv_a_proj" in n for n in buf_names) + + cm, alloc = _make_latent_cache_manager(cfg_absorb, torch.bfloat16) + try: + prefill_logits = _step(submodule, cm, "prefill", prompt) + assert prefill_logits.shape == (1, cfg_absorb.vocab_size) + assert torch.isfinite(prefill_logits).all() + # Absorbed == naive up to bf16 rounding through a 2-layer stack + real backend. + torch.testing.assert_close(prefill_logits, ref_logits, rtol=5e-2, atol=5e-2) + + # A few decode steps over the accumulating paged LATENT cache. + next_token = prefill_logits.argmax(-1) + generated = [int(next_token.item())] + for _ in range(4): + logits = _step(submodule, cm, "decode", next_token) + assert logits.shape == (1, cfg_absorb.vocab_size) + assert torch.isfinite(logits).all() + next_token = logits.argmax(-1) + tok = int(next_token.item()) + assert 0 <= tok < cfg_absorb.vocab_size + generated.append(tok) + finally: + alloc.cleanup() + + assert len(generated) == 5 diff --git a/test/modular/test_kimi_mla_absorb.py b/test/modular/test_kimi_mla_absorb.py new file mode 100644 index 000000000..ff8b98201 --- /dev/null +++ b/test/modular/test_kimi_mla_absorb.py @@ -0,0 +1,218 @@ +"""Phase-A confirmation gate for weight-absorbed MLA (``config.mla_absorb``). + +These are CPU-only, GPU-free tests that prove the absorption *math* on the reduced +config, behind the default-off flag. They deliberately do NOT call +``KimiMLAAttention.forward`` (its RMSNorm uses a FlashInfer GPU kernel); instead +they use pure-torch references (mirroring ``test/integration/test_kimi_mla*.py``) +and exercise the real load-time build (``process_weights_after_loading``) + the +absorption algebra + the latent KV-cache config. + +What the absorbed path must satisfy: + * ``w_kc``/``w_vc`` split out of ``kv_b_proj`` reconstruct the naive per-head + ``k_nope``/``v`` from the latent (``test_absorb_reconstructs_kv_b_proj``); + * the absorbed forward math equals the canonical DeepSeek MLA output — the same + thing the naive path reproduces (``test_absorbed_math_matches_deepseek``); + * ``get_kv_cache_config`` reports the shrunk latent cache when the flag is on, + and is byte-identical to naive when off (``test_kv_cache_config_*``). + +The real FlashInfer MLA kernel / paged latent cache (Phase B) is out of scope here +(it is hard-locked to ckv=512/kpe=64, so it cannot run at the reduced dims); the +wired ``forward`` absorbed branch is exercised on GPU in +``test/integration/test_kimi_mla_absorb_forward.py``. + +Run: pytest test/modular/test_kimi_mla_absorb.py -v +""" +import torch +import torch.nn.functional as F + +from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention +from mstar.model.kimi_k2_7.components.rope import ( + _yarn_find_correction_range, + _yarn_linear_ramp_mask, + rotate_gptj, + yarn_get_mscale, +) +from mstar.model.kimi_k2_7.config import KimiK2Config +from mstar.model.kimi_k2_7.kimi_model import KimiK2Model + +# -------------------------------------------------------------------------- +# Pure-torch references (CPU; copied from test_kimi_mla_paged.py so this stays +# self-contained and GPU-free). +# -------------------------------------------------------------------------- + +def _ref_rmsnorm(x, weight, eps): + x32 = x.float() + x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) + return weight * x32.to(x.dtype) + + +def _ref_yarn_rope(pos, q_pe, k_pe, cfg): + r = cfg.rope_scaling + rotary_dim, base, factor = cfg.qk_rope_head_dim, cfg.rope_theta, r["factor"] + max_pos = r["original_max_position_embeddings"] + beta_fast, beta_slow = r.get("beta_fast", 32), r.get("beta_slow", 1) + mscale, mscale_all_dim = r.get("mscale", 1.0), r.get("mscale_all_dim", 0.0) + pos_freqs = base ** (torch.arange(0, rotary_dim, 2).float() / rotary_dim) + ext, interp = 1.0 / pos_freqs, 1.0 / (factor * pos_freqs) + low, high = _yarn_find_correction_range(beta_fast, beta_slow, rotary_dim, base, max_pos) + mask = 1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float) + inv_freq = interp * (1 - mask) + ext * mask + amp = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim) + freqs = torch.outer(pos.float(), inv_freq) + cos = (freqs.cos() * amp).repeat_interleave(2, -1).unsqueeze(-2) + sin = (freqs.sin() * amp).repeat_interleave(2, -1).unsqueeze(-2) + qr = q_pe.float() * cos + rotate_gptj(q_pe.float()) * sin + kr = k_pe.float() * cos + rotate_gptj(k_pe.float()) * sin + return qr.to(q_pe.dtype), kr.to(k_pe.dtype) + + +def _sdpa_causal(q, k, v, scale): + qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) # (H,T,D) + t = q.shape[0] + causal = torch.triu(torch.full((t, t), float("-inf")), diagonal=1) + attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1) + return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) + + +def _q_and_latent(attn, cfg, h, pos): + """Shared Q + normed-latent + roped pe slices (the piece both refs need).""" + t, heads = h.shape[0], attn.num_heads + d_nope, d_rope, latent = cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.kv_lora_rank + eps = cfg.rms_norm_eps + q = _ref_rmsnorm(F.linear(h, attn.q_a_proj.weight), attn.q_a_layernorm.weight, eps) + q = F.linear(q, attn.q_b_proj.weight).view(t, heads, cfg.qk_head_dim) + q_nope, q_pe = q.split([d_nope, d_rope], dim=-1) + lat = F.linear(h, attn.kv_a_proj_with_mqa.weight) + kv_a, k_pe = lat.split([latent, d_rope], dim=-1) + kv_c = _ref_rmsnorm(kv_a, attn.kv_a_layernorm.weight, eps) # (T, L) + k_pe = k_pe.view(t, 1, d_rope) + q_pe, k_pe = _ref_yarn_rope(pos, q_pe, k_pe, cfg) + return q_nope, q_pe, kv_c, k_pe + + +def _deepseek_scale(cfg): + r = cfg.rope_scaling + mscale = yarn_get_mscale(r["factor"], r.get("mscale_all_dim", 0.0)) + return cfg.qk_head_dim ** -0.5 * mscale * mscale + + +def _ref_deepseek_mla(attn, cfg, h, pos): + """Canonical (naive) DeepSeek MLA: materialize k_nope/v, SDPA at Dqk, o_proj.""" + t, heads = h.shape[0], attn.num_heads + d_nope, d_rope, d_v = cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim + q_nope, q_pe, kv_c, k_pe = _q_and_latent(attn, cfg, h, pos) + kv = F.linear(kv_c, attn.kv_b_proj.weight).view(t, heads, d_nope + d_v) + k_nope, v = kv.split([d_nope, d_v], dim=-1) + q = torch.cat([q_nope, q_pe], dim=-1) + k = torch.cat([k_nope, k_pe.expand(t, heads, d_rope)], dim=-1) + out = _sdpa_causal(q, k, v, _deepseek_scale(cfg)).reshape(t, heads * d_v) + return F.linear(out, attn.o_proj.weight) + + +def _absorbed_mla(attn, cfg, h, pos): + """Weight-absorbed MLA: fold w_kc into q, MQA over the latent, fold w_vc into o.""" + t, heads = h.shape[0], attn.num_heads + d_rope, d_v, latent = cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank + q_nope, q_pe, kv_c, k_pe = _q_and_latent(attn, cfg, h, pos) + q_nope = torch.einsum("thd,hdl->thl", q_nope, attn.w_kc) # (T,H,L) + query = torch.cat([q_nope, q_pe], dim=-1) # (T,H,L+Drope) + kv_c_h = kv_c.unsqueeze(1).expand(t, heads, latent) # MQA: shared over heads + key = torch.cat([kv_c_h, k_pe.expand(t, heads, d_rope)], dim=-1) + attn_latent = _sdpa_causal(query, key, kv_c_h, _deepseek_scale(cfg)) # (T,H,L) + out = torch.einsum("thl,hdl->thd", attn_latent, attn.w_vc).reshape(t, heads * d_v) + return F.linear(out, attn.o_proj.weight) + + +def _build_attention_cpu(seed=0): + """Reduced-config KimiMLAAttention on CPU (fp32), absorbed weights built.""" + torch.manual_seed(seed) + cfg = KimiK2Config.reduced() + cfg.mla_absorb = True + attn = KimiMLAAttention(cfg) # CPU, float32 + for lin in (attn.q_a_proj, attn.q_b_proj, attn.kv_a_proj_with_mqa, + attn.kv_b_proj, attn.o_proj): + lin.weight.data.normal_(0, 0.03) + for norm in (attn.q_a_layernorm, attn.kv_a_layernorm): + norm.weight.data.normal_(1.0, 0.02) + attn.process_weights_after_loading() # split kv_b_proj -> w_kc / w_vc + return attn, cfg + + +# -------------------------------------------------------------------------- +# Tests +# -------------------------------------------------------------------------- + +def test_absorb_reconstructs_kv_b_proj(): + """w_kc/w_vc built from kv_b_proj reproduce naive per-head k_nope / v.""" + attn, cfg = _build_attention_cpu(seed=1) + heads, d_nope, d_v, latent = ( + attn.num_heads, cfg.qk_nope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank) + assert tuple(attn.w_kc.shape) == (heads, d_nope, latent) + assert tuple(attn.w_vc.shape) == (heads, d_v, latent) + + kv_c = torch.randn(5, latent) + kv = F.linear(kv_c, attn.kv_b_proj.weight).view(5, heads, d_nope + d_v) + k_nope_ref, v_ref = kv.split([d_nope, d_v], dim=-1) + + k_nope_abs = torch.einsum("tl,hdl->thd", kv_c, attn.w_kc) + v_abs = torch.einsum("tl,hdl->thd", kv_c, attn.w_vc) + torch.testing.assert_close(k_nope_abs, k_nope_ref, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(v_abs, v_ref, rtol=1e-5, atol=1e-5) + + +def test_fused_qkv_a_proj(): + """The fused latent down-projection buffer == cat(q_a_proj, kv_a_proj_with_mqa) + exactly (it is the one GEMM the absorbed forward runs before splitting).""" + attn, cfg = _build_attention_cpu(seed=3) + expected = torch.cat( + [attn.q_a_proj.weight, attn.kv_a_proj_with_mqa.weight], dim=0) + assert tuple(attn.fused_qkv_a_proj_weight.shape) == ( + cfg.q_lora_rank + cfg.kv_lora_rank + cfg.qk_rope_head_dim, cfg.hidden_size) + # Byte-for-byte concat; no arithmetic, so exact equality (rtol/atol 0). + torch.testing.assert_close( + attn.fused_qkv_a_proj_weight, expected, rtol=0, atol=0) + + +def test_absorbed_math_matches_deepseek(): + """The absorbed forward math == the canonical DeepSeek MLA output (the naive + invariant). This is the Phase-A algorithm gate.""" + attn, cfg = _build_attention_cpu(seed=2) + t = 7 + h = torch.randn(t, cfg.hidden_size) * 0.1 + pos = torch.arange(t) + + absorbed = _absorbed_mla(attn, cfg, h, pos) + reference = _ref_deepseek_mla(attn, cfg, h, pos) + + assert absorbed.shape == (t, cfg.hidden_size) + # Pure fp32 algebra; the only difference is float op ordering (two bmms + + # latent SDPA vs materialized SDPA), so the residual is tiny. + torch.testing.assert_close(absorbed, reference, rtol=1e-4, atol=1e-4) + + +def _kv_cfg(mla_absorb): + m = object.__new__(KimiK2Model) # skip __init__ (tokenizer/weights/GPU) + m.config = KimiK2Config.reduced() + m.config.mla_absorb = mla_absorb + return m.get_kv_cache_config()[0] + + +def test_kv_cache_config_absorbed_shrinks_latent(): + cfg = KimiK2Config.reduced() + kv = _kv_cfg(mla_absorb=True) + assert kv.num_kv_heads == 1 + assert kv.head_dim == cfg.kv_lora_rank + cfg.qk_rope_head_dim # 32 + 8 = 40 + assert kv.num_qo_heads == cfg.num_attention_heads # 4 (q still sharded) + assert kv.attention_backend == "mla_absorb" + # per-token cache shrink vs naive padded MHA: 2 * 4 * 64 = 512 -> 1 * 40 = 40 + naive_elems = 2 * cfg.num_attention_heads * cfg.padded_head_dim + absorbed_elems = kv.num_kv_heads * kv.head_dim + assert naive_elems == 512 and absorbed_elems == 40 + + +def test_kv_cache_config_flag_off_is_naive(): + cfg = KimiK2Config.reduced() + kv = _kv_cfg(mla_absorb=False) + assert kv.num_kv_heads == cfg.num_attention_heads # 4 + assert kv.head_dim == cfg.padded_head_dim # 64 + assert kv.attention_backend == "flashinfer" # default naive backend From 260724e43daee26cb26d63064a046fb3158ddf01 Mon Sep 17 00:00:00 2001 From: Garv Ghai Date: Sat, 1 Aug 2026 13:09:38 +0000 Subject: [PATCH 6/9] Kimi-K2.7: trim verbose comments and docstrings Reduce comment density across the Kimi-K2.7 port and its supporting engine/utils changes to match the surrounding mstar style. Explanatory prose that duplicated the code, reference-file line citations, and narrative docstrings are condensed to short summaries; the non-obvious rationale (meta/to_empty buffer handling, MLA softmax scale, cache layout) is kept as one- or two-line notes. Comments only -- no behavioral change. Kimi CPU tests (27) unchanged and ruff reports no new findings. --- configs/kimi_k2_7.yaml | 6 +- configs/kimi_k2_7_code_tp8.yaml | 34 +--- configs/kimi_k2_7_code_tp8_shm.yaml | 7 +- configs/kimi_k2_7_repro.yaml | 27 +-- configs/kimi_k2_7_tp2.yaml | 26 +-- mstar/distributed/communication.py | 7 +- mstar/engine/cache_manager.py | 118 ++--------- mstar/engine/cuda_graph_runner.py | 7 +- mstar/engine/kv_cache_engine.py | 7 +- mstar/engine/kv_store.py | 13 +- .../model/components/quantization/__init__.py | 8 +- mstar/model/components/quantization/base.py | 28 +-- .../components/quantization/marlin_moe.py | 28 +-- mstar/model/kimi_k2_7/_testing.py | 15 +- mstar/model/kimi_k2_7/components/attention.py | 155 +------------- mstar/model/kimi_k2_7/components/causal_lm.py | 32 +-- .../kimi_k2_7/components/decoder_layer.py | 30 +-- .../kimi_k2_7/components/language_model.py | 38 +--- mstar/model/kimi_k2_7/components/moe.py | 192 +----------------- mstar/model/kimi_k2_7/components/rope.py | 28 +-- mstar/model/kimi_k2_7/config.py | 183 ++--------------- mstar/model/kimi_k2_7/kimi_model.py | 162 +-------------- mstar/model/kimi_k2_7/quantization.py | 107 +--------- mstar/model/kimi_k2_7/submodules.py | 75 +------ mstar/model/kimi_k2_7/weight_loader.py | 103 +--------- mstar/utils/flashinfer_utils.py | 68 +------ mstar/utils/fused_moe/kernels.py | 58 ++---- mstar/utils/fused_moe/runner.py | 17 +- mstar/utils/marlin/__init__.py | 8 +- mstar/utils/marlin/loader.py | 26 +-- mstar/utils/marlin/ops.py | 47 +---- mstar/utils/marlin/scalar_type.py | 15 +- pyproject.toml | 2 - test/integration/test_kimi_components.py | 31 +-- test/integration/test_kimi_decoder_layer.py | 38 +--- .../test_kimi_flashinfer_attention.py | 61 +----- test/integration/test_kimi_forward.py | 45 +--- test/integration/test_kimi_mla.py | 32 +-- .../test_kimi_mla_absorb_forward.py | 23 --- .../test_kimi_mla_absorb_kernel.py | 99 ++------- .../test_kimi_mla_absorb_marlin_merge.py | 29 --- .../integration/test_kimi_mla_absorb_paged.py | 45 +--- .../integration/test_kimi_mla_absorb_serve.py | 46 ----- test/integration/test_kimi_mla_paged.py | 48 +---- test/integration/test_kimi_moe.py | 46 +---- .../test_kimi_moe_inkernel_dequant.py | 38 +--- test/integration/test_kimi_moe_marlin.py | 26 +-- ...test_kimi_quant_inkernel_weight_loading.py | 92 +-------- .../test_kimi_quant_weight_loading.py | 74 +------ test/integration/test_kimi_serve_e2e.py | 73 +------ test/integration/test_kimi_submodule.py | 69 +------ test/integration/test_kimi_tp.py | 92 +-------- test/integration/test_kimi_weight_loading.py | 67 +----- test/integration/test_marlin_kernels.py | 23 +-- test/modular/test_kimi_k27_code_wiring.py | 52 +---- test/modular/test_kimi_mla_absorb.py | 46 +---- test/modular/test_kimi_model.py | 18 +- test/modular/test_kimi_quant.py | 96 +-------- 58 files changed, 195 insertions(+), 2791 deletions(-) diff --git a/configs/kimi_k2_7.yaml b/configs/kimi_k2_7.yaml index 7291213b9..c4cac0e7a 100644 --- a/configs/kimi_k2_7.yaml +++ b/configs/kimi_k2_7.yaml @@ -1,9 +1,5 @@ model: "kimi_k2_7" -# Kimi-K2.7 text backbone (DeepSeek-V3). This single-GPU config (tp_size 1) is the -# bring-up / reduced-model deployment. The real 1T-param MoE needs TP8 across a -# node — set `ranks: [0,1,2,3,4,5,6,7]` and `tp_size: 8` — and, being >1 GPU, the -# single-node tensor transport TENSOR_PROTOCOL=SHM (or TCP; never RDMA on -# coriander). One KV_CACHE LLM node with the prefill + decode(Loop) walks. +# Single-rank Kimi-K2.7 text config; use TP8 configs for the real 1T checkpoint. max_seq_len: 262144 node_groups: - node_names: [LLM] diff --git a/configs/kimi_k2_7_code_tp8.yaml b/configs/kimi_k2_7_code_tp8.yaml index 35ee7376d..3f9d48aeb 100644 --- a/configs/kimi_k2_7_code_tp8.yaml +++ b/configs/kimi_k2_7_code_tp8.yaml @@ -1,45 +1,13 @@ model: "kimi_k2_7" -# FALLBACK ONLY — prefer kimi_k2_7_code_tp8_shm.yaml -# Kimi-K2.7 text backbone — REAL moonshotai/Kimi-K2.7-Code single-node TP=8 serve -# config (INT4, text-only). This points at the actual ~595 GB multimodal -# KimiK25ForConditionalGeneration checkpoint but serves ONLY its DeepSeek-V3 text -# path: the DeepSeek-V3 dims live under `text_config`, its `quantization_config` -# (pack-quantized, num_bits=4, group_size=32, routed experts only) is NESTED under -# `text_config`, and the vision tower (vision_tower.* / mm_projector.*) is dropped -# on load. Serves the routed experts packed with in-kernel Triton W4A16 dequant — -# the only path that fits the real weights. -# -# `model_kwargs` (forwarded to KimiK2Model.__init__): -# * config_variant: k27_code -> KimiK2Config.k27_code() (full 1T dims + -# moe_in_kernel_dequant=True; keeps the default -# beta_fast=32.0 — the K2.7-Code text_config value). -# The checkpoint's nested -# `text_config.quantization_config` is auto-read at -# load, so the routed experts are served packed (int32); -# the lm_head / MLA / dense-FFN / shared experts / -# vision stay bf16, matching the checkpoint `ignore`. -# * checkpoint_path -> local HF-format snapshot of the 595 GB download -# * tokenizer_mode: hf -> the real Kimi tokenizer (tiktoken-based; ids >> 256) -# -# max_seq_len is trimmed to 8192 for first bring-up (the real max is 262144); keep -# the paged KV cache modest for a first single-request serve. -# -# On this cluster (coriander) RDMA is unavailable, so launch with a non-RDMA tensor -# transport: TENSOR_PROTOCOL=SHM (single node, all 8 GPUs) — see mstar/.sample.env. +# Fallback TP8 config for local Kimi-K2.7-Code snapshot; prefer the /dev/shm variant. max_seq_len: 8192 model_kwargs: config_variant: k27_code checkpoint_path: /m-coriander/coriander/garv901/kimi_k2_7_code tokenizer_mode: hf -# Paged KV cache sized for the real MLA head dims (num_attention_heads=64 KV heads -# in the naive/materialized MLA path, padded_head_dim=256). 512 pages * 128 = -# 65536-token capacity — modest, for a first single-request bring-up serve. kv_cache: max_num_pages: 512 page_size: 128 -# The LLM node is the tensor-parallel node: all 8 ranks on one node with tp_size 8 -# for both the prefill and decode walks. The model code is TP-agnostic — the -# per-rank shard degree comes from here. node_groups: - node_names: [LLM] ranks: [0, 1, 2, 3, 4, 5, 6, 7] diff --git a/configs/kimi_k2_7_code_tp8_shm.yaml b/configs/kimi_k2_7_code_tp8_shm.yaml index 1b8ec258c..bbd67205c 100644 --- a/configs/kimi_k2_7_code_tp8_shm.yaml +++ b/configs/kimi_k2_7_code_tp8_shm.yaml @@ -1,10 +1,5 @@ model: "kimi_k2_7" -# Attempt-3 variant of kimi_k2_7_code_tp8.yaml: identical EXCEPT the checkpoint is -# loaded from /dev/shm (RAM-backed tmpfs) instead of shared ZFS. Attempt 2 loaded -# from ZFS over ~59 min and a rank was silently SIGKILL'd (likely a shared-FS mmap -# SIGBUS / contention reaper during the long load). Staging the 595GB checkpoint -# into /dev/shm makes the load RAM-speed (minutes) and immune to mmap I/O faults. -# Stage first with tools: /m-coriander/coriander/garv901/stage_to_shm.sh +# TP8 Kimi-K2.7-Code config using a RAM-backed /dev/shm checkpoint copy. max_seq_len: 8192 model_kwargs: config_variant: k27_code diff --git a/configs/kimi_k2_7_repro.yaml b/configs/kimi_k2_7_repro.yaml index 14f0d86e8..358bcd43d 100644 --- a/configs/kimi_k2_7_repro.yaml +++ b/configs/kimi_k2_7_repro.yaml @@ -1,38 +1,13 @@ model: "kimi_k2_7" -# Kimi-K2.7 text backbone — REDUCED / synthetic single-GPU repro config (serve -# bring-up). It drives the full serving path (API server -> conductor -# -> worker -> KV_CACHE engine -> decode loop -> tokens) on a tiny model that runs -# without the 1T checkpoint. NOT a real deployment — for that use kimi_k2_7.yaml. -# -# `model_kwargs` (forwarded to KimiK2Model.__init__) redirect this model at a -# local synthetic checkpoint + reduced config + a trivial byte tokenizer, so no -# 1T weights and no real Kimi tokenizer are needed: -# * config_variant: reduced -> KimiK2Config.reduced() (vocab 256, 2 layers) -# * checkpoint_path -> the dir written by -# tools/kimi_goldens/make_repro_checkpoint.py -# * tokenizer_mode: byte -> UTF-8 byte identity tokenizer (ids in [0,256)) -# -# Generate the checkpoint first (writes tools/kimi_goldens/repro/checkpoint, which -# is gitignored): -# python tools/kimi_goldens/make_repro_checkpoint.py -# `checkpoint_path` is relative to the mstar package root — run `mstar-serve` from -# there (the launch harness does). Make it absolute if you launch from elsewhere. +# Reduced synthetic serve config. Generate tools/kimi_goldens/repro/checkpoint first. max_seq_len: 512 model_kwargs: config_variant: reduced checkpoint_path: tools/kimi_goldens/repro/checkpoint tokenizer_mode: byte -# Small paged KV cache — the reduced model needs almost nothing (2 layers, 4 heads, -# head_dim 64). 256 pages * 128 = 32768-token capacity, ~33 MB. Sized to cover -# CUDA-graph decode capture (default batch sizes up to 64, double-buffered), not -# just the single serving request. kv_cache: max_num_pages: 256 page_size: 128 -# The CUDA-graph prefill capture grid is trimmed to a single short-prompt bucket -# at batch size 1 by KimiK2Config.reduced() (its prefill_token_buckets / -# prefill_capture_batch_sizes), which `config_variant: reduced` above selects — so -# no env vars are needed. The full 6x5 grid is only captured for the full model. node_groups: - node_names: [LLM] ranks: [0] diff --git a/configs/kimi_k2_7_tp2.yaml b/configs/kimi_k2_7_tp2.yaml index 1bbf87e5c..01ae87549 100644 --- a/configs/kimi_k2_7_tp2.yaml +++ b/configs/kimi_k2_7_tp2.yaml @@ -1,26 +1,5 @@ model: "kimi_k2_7" -# Kimi-K2.7 text backbone — REDUCED / synthetic TP=2 repro config (TP -# correctness). Same tiny synthetic model as kimi_k2_7_repro.yaml, but the LLM -# node runs tensor-parallel across 2 ranks so the MLA head-shard + MoE -# intermediate-shard paths execute end-to-end. NOT a real deployment — the 1T -# model needs TP8 / multi-node (see kimi_k2_7.yaml). -# -# `model_kwargs` (forwarded to KimiK2Model.__init__) redirect this model at a -# local synthetic checkpoint + reduced config + a trivial byte tokenizer: -# * config_variant: reduced -> KimiK2Config.reduced() (vocab 256, 2 layers, -# 4 attention heads -> 2/rank, moe_inter 64 -> 32/rank) -# * checkpoint_path -> the dir written by -# tools/kimi_goldens/make_repro_checkpoint.py -# * tokenizer_mode: byte -> UTF-8 byte identity tokenizer (ids in [0,256)) -# -# Generate the checkpoint first (writes tools/kimi_goldens/repro/checkpoint, which -# is gitignored): -# python tools/kimi_goldens/make_repro_checkpoint.py -# `checkpoint_path` is relative to the mstar package root — run `mstar-serve` from -# there. Make it absolute if you launch from elsewhere. -# -# On this cluster (coriander) RDMA is unavailable, so launch with a non-RDMA -# tensor transport, e.g. TENSOR_PROTOCOL=SHM (single node) — see mstar/.sample.env. +# Reduced synthetic TP=2 config. Generate tools/kimi_goldens/repro/checkpoint first. max_seq_len: 512 model_kwargs: config_variant: reduced @@ -29,9 +8,6 @@ model_kwargs: kv_cache: max_num_pages: 256 page_size: 128 -# The LLM node is the tensor-parallel node (mirrors configs/orpheus_tp2.yaml): -# ranks [0, 1] with tp_size 2 for both the prefill and decode walks. The model -# code is TP-agnostic — the per-rank shard degree comes from here. node_groups: - node_names: [LLM] ranks: [0, 1] diff --git a/mstar/distributed/communication.py b/mstar/distributed/communication.py index f2f7101b2..3fd67061a 100644 --- a/mstar/distributed/communication.py +++ b/mstar/distributed/communication.py @@ -234,9 +234,7 @@ def init_dist( if not self.any_parallelism: return - # A generous timeout (default is ~10 min). First-time bring-up of a very - # large checkpoint — slow weight load, first-ever kernel JIT, and CUDA-graph - # capture can leave ranks waiting at this setup barrier past the default. + # Large-checkpoint load/JIT/capture can exceed the default NCCL timeout. dist.init_process_group( backend="nccl", init_method=init_method, @@ -252,8 +250,6 @@ def init_dist( # an SP group (degenerate meshes) maps to one subgroup. rank_tuple_to_pg: dict[tuple[int, ...], "dist.ProcessGroup"] = {} for rank_tuple in self.world_parallel_groups: - # Same generous timeout as init_process_group above (slow first load, - # kernel JIT, and graph capture can stall ranks past the default). rank_tuple_to_pg[rank_tuple] = dist.new_group( ranks=list(rank_tuple), timeout=timedelta(hours=2) ) @@ -381,4 +377,3 @@ def __init__( self.per_worker_config[worker_ids[rank]].add_sp( node, self.sp_comm_groups[key] ) - diff --git a/mstar/engine/cache_manager.py b/mstar/engine/cache_manager.py index 33d66a978..209fb1525 100644 --- a/mstar/engine/cache_manager.py +++ b/mstar/engine/cache_manager.py @@ -127,11 +127,7 @@ class _PlanState: # segment over its contiguous frozen prefix. None on paged plans, which # keep the FlashInfer path. See DenseGenCacheManager._build_dense_gen_plan. dense_gen: dict | None = None - # Set when MlaAbsorbCacheManager planned this label: the compressed-latent - # scatter indices (token_to_page/token_to_cache) plus a per-request list of - # (q_start, seq_len, total_len, page_indices) that run_attention_mla uses to - # gather each request's full cached latent and run its causal SDPA. None on - # the standard FlashInfer/dense plans. See MlaAbsorbCacheManager.plan_attention. + # MLA absorb fallback plan: latent scatter indices and per-request gather layout. mla: dict | None = None @@ -1524,32 +1520,10 @@ def _mla_kernel_available(ckv: int, kpe: int, sm_major: int) -> bool: class MlaAbsorbCacheManager(FlashInferCacheManager): - """Compressed-latent (weight-absorbed) MLA attention backend. - - Serves DeepSeek/Kimi MLA over a *latent* paged cache: instead of the - standard 6D ``[layers, pages, 2, page_size, kv_heads, head_dim]`` K/V cache, - the KV cache is the 4D latent tensor - ``[layers, pages, page_size, latent_width]`` allocated by - ``KVCacheEngine.load_model`` when ``attention_backend == "mla_absorb"`` - (``latent_width = kv_lora_rank + qk_rope_head_dim``). One compressed latent - vector ``cat([kv_c, k_pe])`` is stored per token (MQA: a single shared latent - "head"), which is ~L/head_dim smaller than materialized K/V. - - The standard FlashInfer paged wrappers assume the 6D layout, so this backend - does not build one. When the FlashInfer **MLA** kernel is available at the real - Kimi dims (see ``_mla_kernel_available``), ``plan_attention`` builds a - ``FlashInferMLAWrapper`` (the dedicated ckv=512/kpe=64 kernel) and - ``run_attention_mla`` scatters the new latents then runs it over strided - ``ckv``/``kpe`` views of the combined cache — this is the production fast path - and is CUDA-graph capturable. Otherwise (reduced configs, pre-sm90, no - flashinfer) it falls back to an all-dims **SDPA** path: ``plan_attention`` - records the per-token write (page, offset) indices — the same index math the - FlashInfer prefill wrapper's ``plan`` computes — plus the per-request gather - layout; ``run_attention_mla`` scatters the new latents, gathers each request's - full cached latent, and runs a causal SDPA in which ``value`` is the first - ``L`` dims of the same latent that forms ``key`` (weight absorption folds - ``k_nope``/``v`` into the query/output projections, so the cache only holds - ``kv_c`` + rope). The SDPA path is eager-only. + """Paged-cache backend for weight-absorbed MLA. + + Uses a 4D latent cache ``[layers, pages, page_size, ckv + kpe]``. Real Kimi + dims on sm90 use FlashInfer MLA; other dims fall back to eager SDPA. """ def plan_attention( @@ -1561,33 +1535,11 @@ def plan_attention( label: str | None = None, **kwargs, ): - """Allocate pages and record the latent attention plan. - - Allocates enough pages for ``seq_len + new_tokens`` per request, then plans - one of two paths (chosen by ``_mla_kernel_available``): - - - **Kernel fast path** (real dims + sm90 + flashinfer): build the MLA index - tensors (qo_indptr / kv_indptr / kv_indices / kv_len_arr) — the same - batch-level tensors ``FlashInferCacheManager._plan_attention_impl`` builds - — and plan a persistent (CUDA-graph) or fresh (eager) ``FlashInferMLAWrapper`` - (stored on ``ps.wrapper``). ``ps.mla`` is cleared; the wrapper owns the - scatter indices. - - **SDPA fallback** (all dims, eager): record, for every new token, the - (page, within-page offset) it will be written to — computed exactly as - ``FlashInferPrefillWrapper.plan`` does (absolute position ``g`` maps to - page ``page_indices[g // page_size]`` at offset ``g % page_size``) — plus - the per-request gather layout, stashed on ``ps.mla``. ``ps.wrapper`` stays - None. - - Planning hints for other backends (``**kwargs``) are ignored. - """ + """Allocate pages and plan the FlashInfer MLA or eager SDPA path.""" self._batched_cfg_info = None effective_label = label if label is not None else self._active_label() - # This backend always re-plans (it does not implement the plan-overlap - # short-circuit); clear any pre-plan marker so it never causes a stale skip. - # The re-plan writes the (per-slot) wrapper's static buffers correctly under - # capture — only the overlap perf optimization is forgone (a follow-up). + # Always re-plan; this backend does not support the overlap skip yet. self._pre_planned_labels.discard(effective_label) if effective_label not in self._plan_states: self._plan_states[effective_label] = _PlanState() @@ -1601,7 +1553,6 @@ def plan_attention( use_kernel = ckv is not None and _mla_kernel_available(ckv, kpe, sm_major) if use_kernel: - # ---- Kernel fast path: build batched MLA index tensors + plan wrapper. qo_indptr_list = [0] kv_indptr_list = [0] all_page_indices: list[int] = [] @@ -1625,8 +1576,7 @@ def plan_attention( if dtype is None: dtype = self.kv_cache.dtype if ps.wrapper is None: - # Eager mode: fresh wrapper each forward (the manager is rebuilt per - # forward). CUDA-graph mode passes a persistent wrapper in ps.wrapper. + # Eager gets a fresh wrapper; CUDA graph injects a persistent one. ps.wrapper = FlashInferMLAWrapper( workspace_buffer=self.buffer_manager.get(effective_label), num_heads=cfg.num_qo_heads, @@ -1647,7 +1597,6 @@ def plan_attention( ) ps.mla = None else: - # ---- SDPA fallback: per-token scatter + per-request gather layout. token_to_page: list[int] = [] token_to_cache: list[int] = [] requests: list[dict] = [] @@ -1663,10 +1612,7 @@ def plan_attention( ) page_indices = state.page_indices - # WRITE indices for the sl new tokens: token j lands at absolute - # position g = old_len + j -> page page_indices[g // page_size], - # offset g % page_size. (Same math as the FlashInfer plan's - # token_to_page / token_to_cache, specialized to per-request order.) + # Same page/offset mapping as FlashInfer's token_to_page plan. for j in range(sl): g = old_len + j token_to_page.append(page_indices[g // page_size]) @@ -1692,8 +1638,7 @@ def plan_attention( "requests": requests, } - # Record like the base plan does so advance_seq_lens / flush_to_store - # see this label's per-request new-token counts and store policy. + # Keep advance_seq_lens / flush_to_store aligned with the base manager. ps.seq_lens = seq_lens ps.write_store = write_store ps.dense_gen = None @@ -1717,14 +1662,6 @@ def run_attention_mla( layer_idx: transformer layer; defaults to self.layer_idx. Returns: [T, H, L] attention output (the ``value`` = ``kv_c`` slice width). - - Writes ``cat([kv_c, k_pe])`` (width L+Drope) as one latent per token into - this layer's paged latent cache at the planned (page, offset) locations, - then attends. When ``plan_attention`` selected the FlashInfer MLA kernel - (``ps.wrapper`` set), the kernel runs over strided ``ckv``/``kpe`` views of - the combined cache; otherwise a causal SDPA gathers each request's full - cached latent (query ``cat([q_nope, q_pe])``, key = the full latent, value = - its first L dims) at ``kv_cache_config.softmax_scale``. """ if layer_idx is None: layer_idx = self.layer_idx @@ -1733,29 +1670,26 @@ def run_attention_mla( ps = self._plan_states[label] assert self.kv_cache is not None - latent_cache = self.kv_cache[layer_idx] # [max_pages, page_size, L+Drope] - latent = torch.cat([kv_c, k_pe], dim=-1).squeeze(1) # [T, L+Drope] + latent_cache = self.kv_cache[layer_idx] + latent = torch.cat([kv_c, k_pe], dim=-1).squeeze(1) if ps.wrapper is not None: - # ---- FlashInfer MLA kernel fast path (CUDA-graph capturable). ps.wrapper.set_latent(latent_cache, latent) L = q_nope.shape[-1] # ckv width (post-w_kc absorption) ckv_cache = latent_cache[..., :L] kpe_cache = latent_cache[..., L:] return ps.wrapper.run(q_nope, q_pe, ckv_cache, kpe_cache).to(q_nope.dtype) - # ---- SDPA fallback. mla = ps.mla assert mla is not None - # Scatter: one latent vector per new token into (page, offset). latent_cache[mla["token_to_page"], mla["token_to_cache"]] = latent.to( latent_cache.dtype ) T, H, L = q_nope.shape scale = self.kv_cache_config.softmax_scale - query_all = torch.cat([q_nope, q_pe], dim=-1) # [T, H, L+Drope] + query_all = torch.cat([q_nope, q_pe], dim=-1) out = torch.empty(T, H, L, dtype=q_nope.dtype, device=q_nope.device) for req in mla["requests"]: @@ -1763,15 +1697,14 @@ def run_attention_mla( sl = req["seq_len"] total_len = req["total_len"] - # Gather this request's full cached latent [total_len, L+Drope] - # (mirror the DenseGenCacheManager page-gather). + # Mirror the dense-gen page gather for the fallback. gathered = latent_cache[req["page_indices"]].reshape( -1, latent_cache.shape[-1] )[:total_len] - key = gathered # [total_len, L+Drope] - value = gathered[:, :L] # [total_len, L] + key = gathered + value = gathered[:, :L] - q_req = query_all[q_start:q_start + sl] # [sl, H, L+Drope] + q_req = query_all[q_start:q_start + sl] out[q_start:q_start + sl] = self._sdpa_mla( q_req, key, value, old_len=total_len - sl, scale=scale ) @@ -1785,18 +1718,11 @@ def _sdpa_mla( old_len: int, scale: float, ) -> torch.Tensor: - """Causal SDPA for one request's MLA step. - - ``q`` is [sl, H, D] (D = L+Drope); ``key`` [total, D] and ``value`` - [total, L] are the single latent "head" broadcast over the H query heads. - Query token j is at absolute position ``old_len + j`` and attends to - cached positions ``0 .. old_len + j`` (causal, includes itself). Handles - both prefill (old_len=0) and a decode step (sl=1, old_len=total-1). - """ + """Causal SDPA fallback for one request over the shared latent head.""" sl = q.shape[0] total = key.shape[0] - qt = q.transpose(0, 1).float() # [H, sl, D] - scores = torch.einsum("hqd,kd->hqk", qt, key.float()) * scale # [H, sl, total] + qt = q.transpose(0, 1).float() + scores = torch.einsum("hqd,kd->hqk", qt, key.float()) * scale q_pos = old_len + torch.arange(sl, device=q.device) k_pos = torch.arange(total, device=q.device) mask = torch.where( @@ -1806,8 +1732,8 @@ def _sdpa_mla( ) scores = scores + mask attn = scores.softmax(-1) - out = torch.einsum("hqk,kd->hqd", attn, value.float()) # [H, sl, L] - return out.transpose(0, 1).to(q.dtype) # [sl, H, L] + out = torch.einsum("hqk,kd->hqd", attn, value.float()) + return out.transpose(0, 1).to(q.dtype) # Backend registry: KVCacheConfig.attention_backend names one of these. diff --git a/mstar/engine/cuda_graph_runner.py b/mstar/engine/cuda_graph_runner.py index 6da5db6d2..7637ad01a 100644 --- a/mstar/engine/cuda_graph_runner.py +++ b/mstar/engine/cuda_graph_runner.py @@ -337,12 +337,7 @@ def _create_persistent_wrappers( cfg = self.kv_cache_config - # Compressed-latent MLA fast path: when the backend is "mla_absorb" and the - # FlashInfer MLA kernel is available at these dims (real Kimi dims on sm90), - # capture uses a persistent FlashInferMLAWrapper for BOTH decode and prefill - # (its run() serves both). This is the only capturable absorbed path — the - # SDPA fallback is eager-only, so reduced-dims / non-kernel absorbed serving - # runs eager (no capture). See MlaAbsorbCacheManager. + # Only the FlashInfer MLA path is capturable; absorbed SDPA stays eager. use_mla_kernel = ( cfg.attention_backend == "mla_absorb" and cfg.mla_ckv_dim is not None diff --git a/mstar/engine/kv_cache_engine.py b/mstar/engine/kv_cache_engine.py index 7671d30b3..920289859 100644 --- a/mstar/engine/kv_cache_engine.py +++ b/mstar/engine/kv_cache_engine.py @@ -229,12 +229,7 @@ def load_model( num_kv_heads = cfg.num_kv_heads if cfg.attention_backend == "mla_absorb": - # Compressed-latent MLA cache: one latent vector per token of - # width head_dim (= kv_lora_rank + qk_rope_head_dim). Drop the - # 2-wide K/V axis (a single latent, not a K/V pair) and the - # num_kv_heads axis (MQA — one shared latent head), giving a 4D - # [num_layers, max_pages, page_size, latent_width] cache that - # MlaAbsorbCacheManager scatters into / gathers from. + # One compressed latent per token: no K/V axis and no KV-head axis. kv_cache = torch.zeros( num_layers, max_num_pages, page_size, head_dim, dtype=kv_cache_type, device=device, diff --git a/mstar/engine/kv_store.py b/mstar/engine/kv_store.py index 465431f25..f9f2c6ff9 100644 --- a/mstar/engine/kv_store.py +++ b/mstar/engine/kv_store.py @@ -122,17 +122,9 @@ class KVCacheConfig: # FA3 on Hopper; models can pin ``fa2`` when their deployment toolchain # cannot compile the Hopper JIT kernels. flashinfer_backend: str = "auto" - # Softmax scale for the compressed-latent MLA backend ("mla_absorb"). MLA's - # intended scale is qk_head_dim**-0.5 * mscale**2, which differs from the - # 1/sqrt(head_dim) a standard kernel would apply over the latent width, so - # the model passes the correct value here and MlaAbsorbCacheManager reads it - # in run_attention_mla. None for the standard paged backends (unused). + # For "mla_absorb", whose scale is based on qk_head_dim rather than latent width. softmax_scale: float | None = None - # For "mla_absorb": the compressed-KV latent width (kv_lora_rank), i.e. the - # ``ckv`` half of the combined latent ``head_dim = ckv + kpe``. The FlashInfer - # MLA kernel fast path needs this split at plan time to pass head_dim_ckv / - # head_dim_kpe. None for the standard paged backends (unused; SDPA fallback - # derives the split from the query shapes at run time). + # For "mla_absorb": split combined latent head_dim into ckv + kpe. mla_ckv_dim: int | None = None def __post_init__(self): @@ -740,4 +732,3 @@ def reload_request(self, request_id: str, cpu_pool) -> None: state.seq_len = seq_len state.position_id_start = pos_id cpu_pool.sync() - diff --git a/mstar/model/components/quantization/__init__.py b/mstar/model/components/quantization/__init__.py index d111578e6..93d1b85dd 100644 --- a/mstar/model/components/quantization/__init__.py +++ b/mstar/model/components/quantization/__init__.py @@ -1,10 +1,4 @@ -"""Model-agnostic quantization backends for mstar. - -Currently provides the Marlin W4A16 routed-expert (fused-MoE) path. The seam -(:class:`FusedMoEQuantizeMethod`) and the generic post-load pass -(:func:`process_weights_after_loading`) are model-agnostic; Kimi-K2.7 is the first -consumer. See :mod:`mstar.model.components.quantization.base`. -""" +"""Model-agnostic quantization backends for mstar.""" from mstar.model.components.quantization.base import ( FusedMoEQuantizeMethod, process_weights_after_loading, diff --git a/mstar/model/components/quantization/base.py b/mstar/model/components/quantization/base.py index 30c4b8e30..1ff017962 100644 --- a/mstar/model/components/quantization/base.py +++ b/mstar/model/components/quantization/base.py @@ -1,23 +1,4 @@ -"""Model-agnostic quantization seams for mstar. - -mstar has no vLLM-style quant-method abstraction. This module introduces the -minimal seam needed to bolt a kernel backend (currently Marlin W4A16 for the -routed experts) onto a model without the model code knowing which kernel runs: - -* :class:`FusedMoEQuantizeMethod` — the interface an MoE block delegates its - quantized routed-expert GEMM to. A block holds one instance, calls - :meth:`~FusedMoEQuantizeMethod.prepare` once post-load to transform its loaded - packed params into the backend's kernel layout, then calls - :meth:`~FusedMoEQuantizeMethod.apply` each forward. Kimi's - :class:`~mstar.model.components.quantization.marlin_moe.MarlinMoEMethod` is the - first implementation; another MoE model can reuse it verbatim. - -* :func:`process_weights_after_loading` — a generic post-load pass. mstar builds - a module on ``meta``, ``to_empty``\\s it, and loads weights, but has no hook to - finalize a kernel layout on the real device afterwards (Marlin needs a one-time - repack + workspace alloc). This walker calls a ``process_weights_after_loading`` - method on every submodule that exposes one; it is a no-op for a plain bf16 model. -""" +"""Model-agnostic quantization hooks for post-load kernel layout fixes.""" from __future__ import annotations from typing import Protocol, runtime_checkable @@ -45,10 +26,6 @@ def prepare( ) -> None: """Transform the loaded compressed-tensors packed expert weights into the backend's runtime layout (e.g. Marlin repack), storing them internally. - - ``w13_packed``/``w2_packed`` are int32 ``(E, N, K // pack_factor)`` and - ``w13_scale``/``w2_scale`` are ``(E, N, K // group_size)`` — the layout - Kimi's Hook B packed params already carry. """ ... @@ -76,8 +53,7 @@ def process_weights_after_loading(root: nn.Module, device: torch.device) -> None Call once after ``load_weights`` and before ``eval()``/CUDA-graph capture. Any submodule exposing a ``process_weights_after_loading(device)`` method gets it invoked (e.g. a Marlin MoE block repacks its packed experts + allocates a - workspace). Modules without the method are skipped, so this is a no-op for a - plain bf16 model. + workspace). """ for module in root.modules(): hook = getattr(module, "process_weights_after_loading", None) diff --git a/mstar/model/components/quantization/marlin_moe.py b/mstar/model/components/quantization/marlin_moe.py index 2400220d8..e7000a296 100644 --- a/mstar/model/components/quantization/marlin_moe.py +++ b/mstar/model/components/quantization/marlin_moe.py @@ -1,14 +1,4 @@ -"""Marlin W4A16 backend for the routed-expert (fused-MoE) GEMM. - -Implements :class:`~mstar.model.components.quantization.base.FusedMoEQuantizeMethod` -for symmetric INT4 (compressed-tensors ``pack-quantized``, group-wise). It consumes -the packed params an MoE block already loaded (the "Hook B" layout Kimi uses — -``(E, N, K // pack_factor)`` int32 weights + ``(E, N, K // group_size)`` bf16 group -scales), repacks them once into Marlin's tiled layout, and thereafter runs the -vendored Marlin kernels (:mod:`mstar.utils.marlin`). - -Model-agnostic: any MoE block with the same packed-param convention can hold one. -""" +"""Marlin W4A16 backend for routed-expert MoE GEMMs.""" from __future__ import annotations import torch @@ -46,15 +36,6 @@ def prepare( w2_scale: torch.Tensor, device: torch.device, ) -> None: - """Repack Hook-B packed experts into Marlin layout. - - mstar packs experts along the input axis as ``(E, N_out, K_in // pack)`` - (compressed-tensors), whereas Marlin's ``gptq_marlin_moe_repack`` wants - GPTQ-style ``(E, K_in // pack, N_out)`` — hence the transpose before each - repack. Scales are permuted the same way. ``w13`` is the fused gate+up - (``N_out = 2 * shard_inter``, ``K_in = hidden``); ``w2`` is the down - projection (``N_out = hidden``, ``K_in = shard_inter``). - """ pf, gs = self.pack_factor, self.group_size E, two_inter, hidden_over_pack = w13_packed.shape hidden = hidden_over_pack * pf @@ -109,12 +90,7 @@ def apply( @staticmethod def shapes_are_legal(hidden: int, shard_inter: int, group_size: int) -> bool: - """Whether the per-rank expert GEMM shapes satisfy Marlin's tile rules. - - Marlin needs ``n % 64 == 0`` and ``k % 128 == 0`` on each GEMM, plus - ``k % group_size == 0``. gate_up: k=hidden, n=2*shard_inter; down: - k=shard_inter, n=hidden. - """ + """Marlin needs n%64 and k%128 for both expert GEMMs.""" if group_size not in (-1, 32, 64, 128): return False checks = [ diff --git a/mstar/model/kimi_k2_7/_testing.py b/mstar/model/kimi_k2_7/_testing.py index ae64329f3..b814f089c 100644 --- a/mstar/model/kimi_k2_7/_testing.py +++ b/mstar/model/kimi_k2_7/_testing.py @@ -26,20 +26,7 @@ def fake_quantize_weight( symmetric: bool = True, scale_dtype: torch.dtype = torch.float32, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Quantize ``weight`` group-wise and return ``(packed, scale, dequant)``. - - A test/harness helper (not used at serve time): produces the on-disk - compressed-tensors tensors *and* the exact bf16 result they dequantize back - to, so a golden can assert the loader reproduces ``dequant`` bit-for-bit. Only - symmetric INT-style quantization is implemented (Kimi's scheme). - - ``dequant`` is derived from the *returned* ``scale`` via :func:`dequantize_weight`, - so it stays consistent with whatever ``scale_dtype`` the scale is stored at. - Pass ``scale_dtype=torch.bfloat16`` to match a real compressed-tensors - checkpoint (whose ``weight_scale`` is stored in the model dtype) — otherwise - the loader's bf16-scale dequant would differ from an fp32-scale reference in - the low bits. - """ + """Return packed weights, stored scales, and the exact dequantized reference.""" if not symmetric: raise NotImplementedError("fake_quantize_weight: only symmetric is implemented") out_f, in_f = weight.shape diff --git a/mstar/model/kimi_k2_7/components/attention.py b/mstar/model/kimi_k2_7/components/attention.py index d611adf57..27100b5b5 100644 --- a/mstar/model/kimi_k2_7/components/attention.py +++ b/mstar/model/kimi_k2_7/components/attention.py @@ -1,58 +1,4 @@ -"""Kimi-K2.7 / DeepSeek-V3 MLA attention — naive path + weight-absorbed path. - -MLA compresses q and k/v through low-rank latents. Two forwards live here, picked -by ``config.mla_absorb`` (default ``True`` -> absorbed): - -* **weight-absorbed** (``mla_absorb=True``, DEFAULT): folds ``kv_b_proj``'s - up-projection into the Q path (``W_UK``) and the O path (``W_UV``) at load (plus - the ``fused_qkv_a_proj`` down-proj fusion), via - :meth:`KimiMLAAttention.process_weights_after_loading`, so attention runs as MQA - over the COMPRESSED latent (``kv_c | k_pe``, one KV head) via - ``cache_handle.run_attention_mla`` — a ~57x per-token KV shrink, numerically - identical to naive up to fp rounding. See :meth:`_forward_absorbed`. Served by - ``engine/cache_manager.py::MlaAbsorbCacheManager`` over a 4D latent paged cache. - That backend currently uses a torch SDPA-over-latent path (correct + memory-lean - but EAGER-ONLY, no CUDA-graph capture); the FlashInfer MLA kernel + CUDA-graph - capture for real-1T throughput is a follow-up. -* **naive / materialized** (``mla_absorb=False``): projects the latent back up to - full per-head K/V and runs ordinary attention, dropping straight onto mstar's - paged ``run_attention`` ``[tokens, heads, head_dim]`` interface (matching vLLM's - ``DeepseekV2Attention``, the non-absorbed class). Zero-pads q/k/v to - ``padded_head_dim`` and folds a softmax boost into q. The M4-golden parity - reference + opt-out fallback (production should keep this until the MLA kernel - lands). - -Per-token shape story (H heads, Dnope=qk_nope, Drope=qk_rope, Dqk=Dnope+Drope, -Dv=v_head_dim, L=kv_lora_rank): - - q: ``q_a_proj`` -> ``q_a_layernorm`` -> ``q_b_proj`` -> ``[T,H,Dqk]``, split - into ``q_nope[..,Dnope]`` / ``q_pe[..,Drope]``. - - kv: ``kv_a_proj_with_mqa`` -> ``[L | Drope]``; the ``L`` slice is RMS-normed - and ``kv_b_proj``-ed to per-head ``[k_nope[..,Dnope] | v[..,Dv]]``; the trailing - ``Drope`` slice is the single shared MQA rope key ``k_pe[T,1,Drope]``. - - YARN RoPE rotates only ``q_pe`` (per head) and ``k_pe`` (broadcast to H heads). - - assemble ``k = [k_nope | k_pe_broadcast] -> [T,H,Dqk]``; zero-pad ``q``/``k`` - (Dqk) and ``v`` (Dv) up to ``padded_head_dim`` (FlashInfer SM90 rejects - ``head_dim_vo`` not in {64,128,256}); fold the scale boost into ``q`` - (``run_attention`` uses the fixed ``1/sqrt(padded_head_dim)`` scale), attend, - slice the output back to ``Dv``, ``o_proj``. - -Cache config for this node: ``num_kv_heads == num_qo_heads == num_attention_heads``, -``head_dim == padded_head_dim`` (256 for the real Dqk=192, 64 for the reduced -Dqk=24). Under tensor parallelism each rank materializes only its -``num_attention_heads // tp_size`` local heads (K/V and Q both shard on the head -axis — there is no separate KV-head group in the naive path), and the paged cache -reports the matching per-rank count. - -The absorbed cache config (``num_kv_heads == 1``, ``head_dim == kv_lora_rank + -qk_rope_head_dim``, no pad) is reported by ``kimi_model.py::get_kv_cache_config`` -when ``mla_absorb`` is set. - -The ``fused_qkv_a_proj`` weight fusion (``q_a_proj`` + ``kv_a_proj_with_mqa`` -> one -GEMM) is applied in the ABSORBED path only: ``process_weights_after_loading`` concats -the two replicated latent down-projection weights into ``fused_qkv_a_proj_weight`` and -:meth:`_forward_absorbed` runs a single ``F.linear`` then splits. The naive ``forward`` -keeps the two separate calls unchanged. -""" +"""Kimi-K2.7 MLA attention with absorbed and naive fallback paths.""" from __future__ import annotations import math @@ -70,26 +16,11 @@ class KimiMLAAttention(nn.Module): - """Multi-head Latent Attention (naive/materialized).""" - def __init__(self, config: KimiK2Config, comm_group: CommGroup | None = None) -> None: super().__init__() if comm_group is None: comm_group = CommGroup.trivial() - # MLA shards on the head dim under TP, mirroring vLLM - # ``DeepseekV2MLAAttention``: the query/kv UP-projections (``q_b_proj`` / - # ``kv_b_proj``) are ColumnParallel and ``o_proj`` is RowParallel, so each - # rank owns a contiguous block of ``num_heads // tp_size`` attention heads. - # The latent DOWN-projections (``q_a_proj`` / ``kv_a_proj_with_mqa``) and - # their RMSNorms are REPLICATED (small shared latent, no head structure). - # ``num_heads`` below is this rank's LOCAL head count used for every - # forward reshape / RoPE / pad / run_attention; the parallel linears are - # given the TOTAL width and divide by tp_size internally, and their - # per-rank ``weight_loader`` slices this rank's head block — so one weight - # path serves tp=1 and tp>1. The paged cache reports the matching per-rank - # head count: ``KVCacheConfig.shard`` divides ``num_qo/kv_heads`` by the - # node's instance world size (tp*sp), exactly like the Orpheus TP path. self.tp_size = comm_group.world_size self.total_num_heads = config.num_attention_heads if self.total_num_heads % self.tp_size != 0: @@ -103,22 +34,14 @@ def __init__(self, config: KimiK2Config, comm_group: CommGroup | None = None) -> self.qk_head_dim = config.qk_head_dim self.v_head_dim = config.v_head_dim self.kv_lora_rank = config.kv_lora_rank - # FlashInfer SM90 rejects head_dim_vo not in {64,128,256}, so q/k/v are - # zero-padded to this width for the paged run_attention; the attention - # output is sliced back to v_head_dim. See config docstring. self.padded_head_dim = config.padded_head_dim - # Parallel linears take the TOTAL head width (they divide by tp_size); - # the forward uses ``self.num_heads`` (local). h = self.total_num_heads - # Q: two-stage low-rank (q_a down -> norm -> q_b up). Down-projections are - # replicated (small rank); up-projections shard over heads under TP. self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=False) self.q_a_layernorm = RMSNorm(config.q_lora_rank, eps=config.rms_norm_eps) self.q_b_proj = ColumnParallelLinear( comm_group, config.q_lora_rank, h * self.qk_head_dim, bias=False) - # KV: shared latent + decoupled rope key. self.kv_a_proj_with_mqa = nn.Linear( config.hidden_size, config.kv_lora_rank + config.qk_rope_head_dim, bias=False) self.kv_a_layernorm = RMSNorm(config.kv_lora_rank, eps=config.rms_norm_eps) @@ -141,36 +64,18 @@ def __init__(self, config: KimiK2Config, comm_group: CommGroup | None = None) -> mscale=rope.get("mscale", 1.0), mscale_all_dim=rope.get("mscale_all_dim", 0.0), ) - # Softmax-scale boost folded into q because run_attention applies a fixed - # 1/sqrt(head_dim) scale and exposes no custom sm_scale. DeepSeek's intended - # softmax scale is ``qk_head_dim**-0.5 * mscale**2``; run_attention now runs - # over the PADDED head dim, so it uses ``padded_head_dim**-0.5``. The - # zero-pad dims contribute 0 to q·k, so to recover the intended scale we - # fold ``mscale**2 * sqrt(padded_head_dim / qk_head_dim)`` into q: - # scores = (q*boost)·k * padded_head_dim**-0.5 - # = q·k * mscale**2 * sqrt(padded/qk) * padded**-0.5 - # = q·k * mscale**2 * qk**-0.5 (the DeepSeek scale). + # run_attention uses 1/sqrt(padded_head_dim); fold DeepSeek's intended + # qk_head_dim**-0.5 * mscale**2 scale into q on the padded path. mscale = yarn_get_mscale(rope["factor"], rope.get("mscale_all_dim", 0.0)) self.softmax_scale_boost = ( mscale * mscale * math.sqrt(self.padded_head_dim / self.qk_head_dim) ) - # DeepSeek's intended softmax scale (uses the PRE-absorption qk_head_dim, - # not the latent width). The naive path reaches it by folding - # softmax_scale_boost into q; the absorbed path hands this to the latent - # cache backend (planned in Phase B; the reduced-config test mock reads it). self.softmax_scale = self.qk_head_dim ** -0.5 * mscale * mscale - # Weight-absorbed MLA (config.mla_absorb): W_UK/W_UV are split out of - # kv_b_proj post-load by process_weights_after_loading(). persistent=False: - # derived tensors, never part of the checkpoint / state_dict. self.mla_absorb = config.mla_absorb if self.mla_absorb: self.register_buffer("w_kc", None, persistent=False) # (H_local, Dnope, L) self.register_buffer("w_vc", None, persistent=False) # (H_local, Dv, L) - # The two replicated latent down-projections (q_a_proj + - # kv_a_proj_with_mqa) both read hidden_states; they fuse into one GEMM - # for the absorbed forward. Built post-load; q_a_proj/kv_a_proj_with_mqa - # remain the checkpoint load targets. persistent=False: derived tensor. self.register_buffer("fused_qkv_a_proj_weight", None, persistent=False) # (q_lora+L+Drope, hidden) def forward( @@ -184,12 +89,10 @@ def forward( num_tokens = hidden_states.shape[0] h = self.num_heads - # --- Q --- q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) q = q.view(num_tokens, h, self.qk_head_dim) q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - # --- KV latent --- latent = self.kv_a_proj_with_mqa(hidden_states) # (T, L + Drope) kv_a, k_pe = latent.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) kv = self.kv_b_proj(self.kv_a_layernorm(kv_a)) @@ -197,23 +100,17 @@ def forward( k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) k_pe = k_pe.view(num_tokens, 1, self.qk_rope_head_dim) # shared MQA rope key - # --- RoPE (only the pe slices) --- q_pe, k_pe = self.rotary(position_ids, q_pe, k_pe) - # --- assemble full q / k (k_pe broadcast over heads) --- q = torch.cat([q_nope, q_pe], dim=-1) # (T, H, Dqk) k_pe = k_pe.expand(num_tokens, h, self.qk_rope_head_dim) k = torch.cat([k_nope, k_pe], dim=-1) # (T, H, Dqk) - # --- zero-pad q/k (Dqk) and v (Dv) up to padded_head_dim for the paged - # run_attention (FlashInfer SM90 requires head_dim_vo in {64,128,256}) --- qk_pad = self.padded_head_dim - self.qk_head_dim q = F.pad(q, [0, qk_pad]) # (T, H, Dpad) k = F.pad(k, [0, qk_pad]) # (T, H, Dpad) v = F.pad(v, [0, self.padded_head_dim - self.v_head_dim]) # (T, H, Dpad) - # --- softmax boost folded into q (compensates padded_head_dim scale), - # attend, strip the pad + v-pad, project --- q = q * self.softmax_scale_boost attn = cache_handle.run_attention(q=q, k=k, v=v) # (T, H, Dpad) attn = attn[..., : self.v_head_dim].reshape(num_tokens, h * self.v_head_dim) @@ -225,16 +122,7 @@ def _forward_absorbed( cache_handle: BatchedCacheManager, position_ids: torch.Tensor, ) -> torch.Tensor: - """Weight-absorbed MLA forward (see module docstring; ``config.mla_absorb``). - - kv_b_proj is folded into Q (``W_UK`` = ``w_kc``) and O (``W_UV`` = ``w_vc``) - so the KV latent stays COMPRESSED and attention is MQA over ``[kv_c | k_pe]`` - (one KV head). Math identity vs naive: ``q_nope · k_nope == - (q_nope @ W_UK) · kv_c`` and ``attn · v == (attn · kv_c) @ W_UV``. The - softmax scale (``self.softmax_scale``) is applied by the latent cache - backend (planned in Phase B; the reduced-config test mock reads it), so — - unlike naive — nothing is folded into q and nothing is padded. - """ + """Run MLA over the compressed latent cache after folding kv_b into Q/O.""" if self.w_kc is None or self.w_vc is None: raise RuntimeError( "mla_absorb forward requires process_weights_after_loading() to " @@ -248,58 +136,33 @@ def _forward_absorbed( num_tokens = hidden_states.shape[0] h = self.num_heads - # --- fused latent down-projection (q_a_proj + kv_a_proj_with_mqa in one - # GEMM), then split into the q latent, the kv latent, and the rope key --- fused = F.linear(hidden_states, self.fused_qkv_a_proj_weight) q_c, kv_a, k_pe = fused.split( [self.q_a_proj.out_features, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - # q_c / kv_a feed the FlashInfer RMSNorm kernel, which requires a - # 64-byte-aligned input pointer. These are mid-tensor split views: q_c is at - # offset 0 (always aligned); kv_a sits at q_lora_rank*dtype bytes (e.g. - # reduced dims -> 96, not 64-aligned). ``.contiguous()`` fixes the multi-token - # (prefill) case but is a NO-OP for a size-1 (decode, T=1) row — a [1,K] view - # is already "contiguous" — leaving kv_a at the misaligned offset. So force a - # fresh contiguous allocation for kv_a via ``clone``. (Real dims align by - # chance; this keeps reduced-config / decode correct too.) + # FlashInfer RMSNorm needs 64-byte input alignment; decode split views can + # be contiguous yet start at an unaligned offset, so clone kv_a. q_c = q_c.contiguous() kv_a = kv_a.clone(memory_format=torch.contiguous_format) - # --- Q (norm -> up), split nope/rope (same as naive) --- q = self.q_b_proj(self.q_a_layernorm(q_c)) q = q.view(num_tokens, h, self.qk_head_dim) q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - # --- KV latent kept COMPRESSED (no kv_b up-projection); norm the kv_c slice --- kv_c = self.kv_a_layernorm(kv_a).view(num_tokens, 1, self.kv_lora_rank) # (T,1,L) k_pe = k_pe.view(num_tokens, 1, self.qk_rope_head_dim) # (T,1,Drope) shared MQA key - # --- RoPE on the pe slices only (q_pe per head, k_pe single shared key) --- q_pe, k_pe = self.rotary(position_ids, q_pe, k_pe) - # --- absorb W_UK into q_nope: (T,H,Dnope) x (H,Dnope,L) -> (T,H,L) --- q_nope = torch.einsum("thd,hdl->thl", q_nope, self.w_kc) - # --- MQA over the latent: key=[kv_c|k_pe], value=kv_c -> (T,H,L) --- attn_latent = cache_handle.run_attention_mla( q_nope=q_nope, q_pe=q_pe, kv_c=kv_c, k_pe=k_pe) - # --- absorb W_UV into the output: (T,H,L) x (H,Dv,L) -> (T,H,Dv) -> o_proj --- out = torch.einsum("thl,hdl->thd", attn_latent, self.w_vc) return self.o_proj(out.reshape(num_tokens, h * self.v_head_dim)) def process_weights_after_loading(self, device: torch.device | str | None = None) -> None: - """Build the absorbed projections from ``kv_b_proj`` (no-op unless mla_absorb). - - Splits ``kv_b_proj.weight`` ``[H_local*(Dnope+Dv), L]`` per local head into - ``w_kc = W_UK [H_local, Dnope, L]`` (absorbed into Q) and - ``w_vc = W_UV [H_local, Dv, L]`` (absorbed into O). Each rank's ``kv_b_proj`` - (ColumnParallelLinear) already holds only its local heads, so this is - TP-correct with no extra sharding — one path serves tp=1 and tp>1. - - Named to match the generic post-load walker protocol so the serve path picks - it up automatically (Phase B). Idempotent. ``device`` is accepted for - protocol compatibility but unused — ``kv_b_proj.weight`` is already resident. - """ + """Build absorbed Q/O projections from the local-head ``kv_b_proj`` shard.""" if not self.mla_absorb: return del device # protocol arg; kv_b_proj.weight already carries the right device @@ -311,9 +174,5 @@ def process_weights_after_loading(self, device: torch.device | str | None = None self.w_kc = w_kc.contiguous() # (H_local, Dnope, L) self.w_vc = w_vc.contiguous() # (H_local, Dv, L) - # Fuse the two replicated latent down-projections into one GEMM weight for - # the absorbed forward: [q_a_proj ; kv_a_proj_with_mqa] -> a single - # (q_lora + kv_lora + Drope, hidden) matmul that splits back into the q - # latent, the kv latent, and the shared rope key. self.fused_qkv_a_proj_weight = torch.cat( [self.q_a_proj.weight, self.kv_a_proj_with_mqa.weight], dim=0).contiguous() diff --git a/mstar/model/kimi_k2_7/components/causal_lm.py b/mstar/model/kimi_k2_7/components/causal_lm.py index 8920927e9..90e190fdd 100644 --- a/mstar/model/kimi_k2_7/components/causal_lm.py +++ b/mstar/model/kimi_k2_7/components/causal_lm.py @@ -1,22 +1,4 @@ -"""Kimi-K2.7 / DeepSeek-V3 assembled text backbone. - -Stacks the :class:`KimiDecoderLayer` blocks between a token embedding and a -final RMSNorm (:class:`KimiLanguageModel`), then wraps that with the untied LM -head (:class:`KimiForCausalLM`). This is the full text forward: token ids → -logits. - -The per-layer cache-handle contract mirrors ``OrpheusLanguageModel`` exactly: -each layer is preceded by ``cache_handle.set_layer_idx(layer_idx)`` (so the paged -KV cache writes/reads the right layer slice), and the loop is followed by a -single ``cache_handle.advance_seq_lens()`` (so every request's ``seq_len`` / -``position_id_start`` steps forward once per forward pass, not once per layer). -The naive MLA reads ``position_ids`` for its YARN RoPE, so unlike Orpheus we -thread ``position_ids`` through each layer. - -Lives in its own module (not ``language_model.py``) to keep the import graph -acyclic: ``decoder_layer`` imports the ``language_model`` builders, so the -assembly that imports ``decoder_layer`` must sit downstream of both. -""" +"""Assembled Kimi-K2.7 text backbone.""" from __future__ import annotations import torch @@ -34,8 +16,6 @@ class KimiLanguageModel(nn.Module): - """Embedding + stacked decoder layers + final norm (returns hidden states).""" - def __init__( self, config: KimiK2Config, comm_group: CommGroup | None = None ) -> None: @@ -66,8 +46,6 @@ def forward( class KimiForCausalLM(nn.Module): - """Text backbone + untied LM head (returns ``[..., vocab]`` logits).""" - def __init__( self, config: KimiK2Config, comm_group: CommGroup | None = None ) -> None: @@ -87,14 +65,6 @@ def forward( return self.lm_head(hidden_states) def load_weights(self, weights, **kwargs) -> set[str]: - """Load an HF DeepSeek-V3 checkpoint stream. - - Called by the shared ``mstar.model.loader.load_weights(model, source, - device)`` driver (mirrors ``OrpheusForCausalLM.load_weights``). Delegates - to :func:`mstar.model.kimi_k2_7.weight_loader.load_kimi_hf_weights` for - the Kimi remap + fused-expert stacked rules. Returns the set of loaded - param paths. - """ from mstar.model.kimi_k2_7.weight_loader import load_kimi_hf_weights packed_experts = ( diff --git a/mstar/model/kimi_k2_7/components/decoder_layer.py b/mstar/model/kimi_k2_7/components/decoder_layer.py index 5048defff..9528f7860 100644 --- a/mstar/model/kimi_k2_7/components/decoder_layer.py +++ b/mstar/model/kimi_k2_7/components/decoder_layer.py @@ -1,23 +1,4 @@ -"""Kimi-K2.7 / DeepSeek-V3 decoder layer. - -One pre-norm transformer block: MLA self-attention then a feed-forward that is -either the dense SwiGLU MLP (the ``first_k_dense_replace`` early layers) or the -fine-grained sigmoid-routed MoE block. Both feed-forwards expose the same -``(x) -> x`` interface, so the residual wiring here is agnostic to which it holds -(``build_mlp_for_layer`` picks per ``layer_idx``). - -This is a Kimi-specific decoder layer rather than the shared -``mstar.model.components.DecoderLayer`` because MLA attention needs -``position_ids`` threaded through its forward (the shared layer's -``self_attn(x, cache_handle=...)`` signature has no position channel — YARN RoPE -is applied inside the attention over the decoupled ``qk_rope`` slice). - -Residual structure mirrors vLLM ``DeepseekV2DecoderLayer.forward``: - residual = h - h = input_layernorm(h); h = self_attn(h, cache, pos); h = residual + h - residual = h - h = post_attention_layernorm(h); h = mlp(h); h = residual + h -""" +"""Kimi-K2.7 decoder layer with MLA position ids threaded through attention.""" from __future__ import annotations import torch @@ -34,15 +15,6 @@ class KimiDecoderLayer(nn.Module): - """Pre-norm MLA + (dense-or-MoE) feed-forward block. - - Args: - config: model config. - layer_idx: index into the stack; selects the dense MLP (``layer_idx < - first_k_dense_replace``) or the MoE block (``build_mlp_for_layer``). - comm_group: TP comm group (trivial single-rank if ``None``). - """ - def __init__( self, config: KimiK2Config, diff --git a/mstar/model/kimi_k2_7/components/language_model.py b/mstar/model/kimi_k2_7/components/language_model.py index cc05ad7e1..a458cfc0d 100644 --- a/mstar/model/kimi_k2_7/components/language_model.py +++ b/mstar/model/kimi_k2_7/components/language_model.py @@ -1,14 +1,4 @@ -"""Kimi-K2.7 language-model builders (DeepSeek-V3 text backbone). - -Thin builders mapping ``KimiK2Config`` onto reused mstar primitives — token -embedding, the dense SwiGLU MLP (the ``first_k_dense_replace`` early layers), -RMSNorm, and the untied LM head — the pieces DeepSeek-V3 shares with a standard -Llama-style stack (the Kimi-specific MLA attention, sigmoid-routed MoE, and YARN -RoPE live elsewhere). Each matches the vLLM DeepSeek-V3 reference -(``deepseek_v2.py``): untied embedding/LM head (``tie_word_embeddings=False``), -``DeepseekV2MLP`` = ``down_proj(SiluAndMul(gate_up_proj(x)))`` (``bias=False``, -silu-only), Llama-style RMSNorm. -""" +"""Kimi-K2.7 language-model builders over existing mstar primitives.""" from __future__ import annotations from mstar.distributed.communication import CommGroup @@ -25,7 +15,6 @@ def build_embedding( config: KimiK2Config, comm_group: CommGroup | None = None ) -> VocabParallelEmbedding: - """Token embedding ``[vocab, hidden]`` (row/vocab-parallel under TP).""" return VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, @@ -37,11 +26,6 @@ def build_embedding( def build_lm_head( config: KimiK2Config, comm_group: CommGroup | None = None ) -> ColumnParallelLinear: - """Untied LM head ``[hidden, vocab]`` (Kimi: ``tie_word_embeddings=False``). - - Column-parallel over vocab with ``gather_output=True`` so the sampler always - sees full ``[..., vocab]`` logits; a no-op all-gather at ``tp_size == 1``. - """ return ColumnParallelLinear( comm_group or CommGroup.trivial(), input_size=config.hidden_size, @@ -52,20 +36,12 @@ def build_lm_head( def build_rmsnorm(config: KimiK2Config) -> RMSNorm: - """Standard Llama-style RMSNorm (not Gemma's ``(1 + weight)`` variant).""" return RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def build_dense_mlp( config: KimiK2Config, comm_group: CommGroup | None = None ) -> ParallelGatedMLP: - """Dense SwiGLU MLP for the ``first_k_dense_replace`` early layers. - - Matches ``DeepseekV2MLP``: fused gate/up projection, ``silu(gate) * up``, - row-parallel down projection, ``bias=False``. Uses the full - ``intermediate_size`` (the MoE layers use ``moe_intermediate_size`` per - expert instead). - """ return ParallelGatedMLP( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, @@ -76,11 +52,6 @@ def build_dense_mlp( def is_moe_layer(config: KimiK2Config, layer_idx: int) -> bool: - """DeepSeek-V3 dense-vs-MoE layer selection. - - The first ``first_k_dense_replace`` layers are dense; thereafter every - ``moe_layer_freq``-th layer is MoE (``deepseek_v2.py`` decoder-layer ctor). - """ return ( layer_idx >= config.first_k_dense_replace and layer_idx % config.moe_layer_freq == 0 @@ -90,19 +61,12 @@ def is_moe_layer(config: KimiK2Config, layer_idx: int) -> bool: def build_moe_block( config: KimiK2Config, comm_group: CommGroup | None = None ) -> KimiSparseMoeBlock: - """Fine-grained MoE block (routed experts + ungated shared expert).""" return KimiSparseMoeBlock(config, comm_group=comm_group) def build_mlp_for_layer( config: KimiK2Config, layer_idx: int, comm_group: CommGroup | None = None ): - """Pick the layer's feed-forward: dense SwiGLU MLP or the MoE block. - - Returns a ``ParallelGatedMLP`` for the early dense layers, else a - ``KimiSparseMoeBlock``. Both expose the same ``(x) -> x`` interface, so the - decoder layer is agnostic to which it holds. - """ if is_moe_layer(config, layer_idx): return build_moe_block(config, comm_group=comm_group) return build_dense_mlp(config, comm_group=comm_group) diff --git a/mstar/model/kimi_k2_7/components/moe.py b/mstar/model/kimi_k2_7/components/moe.py index c7f8e9174..f744f681d 100644 --- a/mstar/model/kimi_k2_7/components/moe.py +++ b/mstar/model/kimi_k2_7/components/moe.py @@ -1,24 +1,4 @@ -"""Kimi-K2.7 / DeepSeek-V3 fine-grained MoE. - -mstar's ``model.components.moe`` router is softmax-only and its shared-expert -block gates the shared expert (Qwen-style). Kimi/DeepSeek-V3 needs a different -router and an *ungated* shared expert, so these live here (append, don't modify -the shared abstraction). The expert dispatch itself is reused verbatim — the -fused-expert GEMM (``fused_experts`` via ``model.components.moe._dispatch``) and -the ``(E, 2*moe_inter, hidden)`` / ``(E, hidden, moe_inter)`` fused param layout. - -Two pieces: - -* :class:`KimiMoEGate` — the router. sigmoid scoring + group-limited top-k - (``n_group`` / ``topk_group``) + ``noaux_tc`` per-expert - ``e_score_correction_bias`` (affects *selection* only; the combine weights come - from the raw sigmoid scores) + optional ``norm_topk_prob`` + a - ``routed_scaling_factor`` folded into the returned weights. Computed in fp32. - Exactly mirrors vLLM ``fused_moe/cpu_fused_moe.py::grouped_topk``. -* :class:`KimiSparseMoeBlock` — router + fused routed experts + ungated shared - expert. ``out = routed(scaled weights) + shared`` (the shared expert does *not* - get ``routed_scaling_factor``). Mirrors vLLM ``deepseek_v2.py::DeepseekV2MoE``. -""" +"""Kimi-K2.7 fine-grained MoE: sigmoid router plus ungated shared expert.""" from __future__ import annotations import logging @@ -39,38 +19,14 @@ logger = logging.getLogger(__name__) -# Log the resolved routed-expert backend once per process (the block is -# instantiated per MoE layer, so a per-block log would repeat ~60x). _BACKEND_LOGGED = False -# --------------------------------------------------------------------------- -# Packed-expert weight loaders (int32 weights + bf16 group scales). -# -# The packed analogs of ``model.components.moe._gate_up_weight_loader`` / -# ``_down_proj_weight_loader`` (which serve the bf16 fused params shared with -# Qwen3-Omni). The TP shard geometry is identical to the bf16 loaders; only the -# last (input/K) axis differs: it is pre-divided by ``pack_factor`` (packed int32) -# or ``group_size`` (bf16 scale). One function serves both the packed and the -# scale tensor for a projection — the divisor is the only difference. -# --------------------------------------------------------------------------- - def _gate_up_packed_loader( tp_rank: int, tp_size: int, full_inter: int, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | int | None = None, ): - """Load one expert's gate_proj/up_proj packed-or-scale tensor into the fused - ``gate_up_proj_packed`` / ``gate_up_proj_scale`` param. - - ``loaded_shard_id`` is ``"gate:N"`` / ``"up:N"``. ``loaded_weight`` is a single - expert's 2-D tensor ``(full_inter, hidden // divisor)`` (divisor = pack_factor - for the int32 packed tensor, group_size for the bf16 scale). The N/out axis - (dim 0) is the TP-sharded one: this rank takes rows - ``[tp_rank*shard_inter : +shard_inter]`` and writes them into the gate half - ``[:shard_inter]`` or up half ``[shard_inter:]`` of ``param[expert]``. The last - axis (the un-sharded input dim) is copied whole. - """ assert loaded_shard_id is not None kind, expert_str = loaded_shard_id.split(":") expert_idx = int(expert_str) @@ -88,17 +44,6 @@ def _down_packed_loader( param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | int | None = None, ): - """Load one expert's down_proj packed-or-scale tensor into ``down_proj_packed`` - / ``down_proj_scale``. - - ``loaded_shard_id`` is ``"down:N"``. ``loaded_weight`` is ``(hidden, moe_inter - // divisor)``; the intermediate dim is the LAST (input) axis and is the - TP-sharded one, already divided by ``divisor`` (pack_factor for packed, - group_size for scale). This rank takes the column stripe - ``[tp_rank*(shard_inter//divisor) : +(shard_inter//divisor)]``. Requires - ``shard_inter % divisor == 0`` (asserted at block build) so the stripe lands on - an int32 / group boundary. - """ assert loaded_shard_id is not None expert_idx = int(str(loaded_shard_id).split(":")[1]) shard_inter = divide(full_inter, tp_size) @@ -108,21 +53,7 @@ def _down_packed_loader( class KimiMoEGate(nn.Module): - """DeepSeek-V3 group-limited sigmoid router with ``noaux_tc`` bias. - - Args: - hidden_size: input hidden dim. - n_routed_experts: number of routed experts (``E``). - num_experts_per_tok: top-k experts per token. - n_group: number of expert groups (``E`` split into ``n_group`` contiguous - groups for group-limited routing). - topk_group: number of groups kept per token. - routed_scaling_factor: scale folded into the returned combine weights. - scoring_func: ``"sigmoid"`` (Kimi/DeepSeek-V3) or ``"softmax"``. - topk_method: ``"noaux_tc"`` enables the per-expert - ``e_score_correction_bias`` (selection-only). Anything else disables it. - norm_topk_prob: renormalize the top-k combine weights to sum to 1. - """ + """DeepSeek-V3 group-limited sigmoid router with selection-only bias.""" def __init__( self, @@ -147,7 +78,6 @@ def __init__( self.topk_method = topk_method self.norm_topk_prob = norm_topk_prob - # Router projection ``[E, hidden]`` (no bias), like DeepSeek ``MoEGate``. self.weight = nn.Parameter(torch.zeros(n_routed_experts, hidden_size)) if topk_method == "noaux_tc": # Per-expert selection bias; fp32, added to scores for group/top-k @@ -161,14 +91,6 @@ def __init__( def forward( self, hidden_states: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: - """Route tokens to experts. - - Returns: - topk_weights: ``(tokens, top_k)`` fp32 combine weights (renormalized - and scaled by ``routed_scaling_factor``). - topk_ids: ``(tokens, top_k)`` int64 expert indices. - """ - # Route in fp32 (DeepSeek runs the router in fp32 for stability). h = hidden_states.reshape(-1, self.hidden_size).float() gating = F.linear(h, self.weight.float()) # (T, E) @@ -223,33 +145,7 @@ def forward( class KimiSparseMoeBlock(nn.Module): - """DeepSeek-V3 MoE block: routed experts + ungated shared expert. - - ``out = routed(x) + shared(x)`` where ``routed`` dispatches the top-k experts - through the fused-expert GEMM with the router's (scaled) combine weights, and - ``shared`` is a plain dense SwiGLU MLP added ungated (no sigmoid gate, no - ``routed_scaling_factor``). - - **TP sharding (intermediate-parallel).** Under tensor parallelism the router - (:class:`KimiMoEGate`) stays REPLICATED — every rank computes the full - ``(top_k_ids, weights)`` — and only the expert GEMMs shard, exactly like - mstar's own ``ParallelSparseMoeBlock``: each rank holds every expert but only - a ``moe_intermediate_size // tp_size`` slice of its SwiGLU intermediate - (``gate_up_proj`` column-parallel, ``down_proj`` row-parallel). The per-rank - partial hidden contributions are summed with a single all-reduce before the - top-k sum-reduce. The shared expert is a ``ParallelGatedMLP`` on the same comm - group, so it shards its intermediate and all-reduces internally. This reuses - the existing fused-expert machinery verbatim and is trivially goldenable - (tp>1 == tp=1). Its tradeoff: every rank still stores ALL experts' weights, so - it does NOT reduce per-rank expert memory — the 1T fit needs true token-dispatch - expert parallelism (all-to-all), which is deliberately not built here. - - Expert weights use the fused layout reused from ``model.components.moe``, - sharded to this rank (``full == moe_intermediate_size``, - ``shard == full // tp_size``): - - ``experts.gate_up_proj``: ``(E, 2 * shard, hidden)`` - - ``experts.down_proj``: ``(E, hidden, shard)`` - """ + """DeepSeek-V3 MoE block: routed experts + ungated shared expert.""" def __init__( self, config: KimiK2Config, comm_group: CommGroup | None = None @@ -263,19 +159,11 @@ def __init__( self.hidden_size = config.hidden_size self.num_experts = config.n_routed_experts self.moe_intermediate_size = config.moe_intermediate_size - # Per-rank slice of each expert's SwiGLU intermediate (== full at tp=1). shard_inter = divide(config.moe_intermediate_size, self.tp_size) - # Packed experts (in-kernel W4A16 dequant) are used iff the checkpoint is - # quantized AND the config opts in. When off, the experts use the bf16 fused - # params (dequantized on load, or a native-bf16 checkpoint loaded directly). self.packed_experts = ( config.quantization_config is not None and config.moe_in_kernel_dequant ) - # Routed-expert W4A16 kernel backend for the packed experts. Marlin layers - # on top of the same packed params; the marlin-vs-triton choice is resolved - # post-load in :meth:`process_weights_after_loading` (a real device is needed - # to probe GPU capability + JIT-build the kernel — ``__init__`` runs on meta). self.quant_kernel = getattr(config, "quant_kernel", "auto") self._marlin_method = None self._use_marlin = False @@ -294,20 +182,11 @@ def __init__( self.experts = nn.Module() if self.packed_experts: - # PACKED expert params (int32 weights + bf16 group scales) INSTEAD of the - # bf16 fused params. Layout mirrors the fused bf16 shapes with the K - # (input) axis compressed: gate_up packs K=hidden, down packs K=inter. - # gate_up_proj_packed: int32 (E, 2*shard_inter, hidden // pack_factor) - # gate_up_proj_scale: bf16 (E, 2*shard_inter, hidden // group_size) - # down_proj_packed: int32 (E, hidden, shard_inter // pack_factor) - # down_proj_scale: bf16 (E, hidden, shard_inter // group_size) qc = config.quantization_config self.group_size = qc.group_size self.pack_factor = qc.pack_factor # 8 for INT4 self.symmetric = qc.symmetric hidden, gs, pf = config.hidden_size, self.group_size, self.pack_factor - # The packed/group axes must divide evenly on BOTH the hidden (gate_up K) - # and the per-rank intermediate stripe (down K, TP-sharded). assert hidden % pf == 0 and hidden % gs == 0, ( f"hidden {hidden} must divide pack_factor {pf} and group_size {gs}" ) @@ -346,18 +225,8 @@ def __init__( shard_inter, ) ) - # The fused expert params are plain nn.Parameters, so they carry no - # per-shard ``weight_loader`` by default. The stacked-param rules route each - # checkpoint expert via a ``"gate:N"/"up:N"/"down:N"`` shard id, so we attach - # the same fused-expert loaders ``ParallelSparseMoeBlock`` uses (or, when - # packed, the packed analogs). The loaders take - # ``(tp_rank, tp_size, full_inter)`` and slice this rank's intermediate - # stripe out of the full-size checkpoint expert — so a single weight path - # serves tp=1 (full) and tp>1 (sharded). self._attach_expert_weight_loaders() - # Ungated shared expert: a dense SwiGLU MLP with the shared intermediate - # size (``moe_intermediate_size * n_shared_experts``). self.shared_expert = ParallelGatedMLP( hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size * config.n_shared_experts, @@ -367,15 +236,7 @@ def __init__( ) def _attach_expert_weight_loaders(self) -> None: - """Give the fused expert params their per-shard ``weight_loader``. - - Mirrors ``ParallelSparseMoeBlock._attach_weight_loaders``. Re-run after - every ``_apply`` (``.to(dtype)`` / ``to_empty(device)`` rebuild the - Parameter objects and drop the attribute), so weights load correctly - through the meta -> to_empty -> load path. When packed, the four packed / - scale params get the Kimi-local packed loaders; otherwise the two bf16 - fused params get the shared loaders. - """ + """Reattach per-shard loaders after ``_apply`` rebuilds parameters.""" from functools import partial full_inter = self.moe_intermediate_size @@ -410,7 +271,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: input_shape = hidden_states.shape flat = hidden_states.view(-1, self.hidden_size).contiguous() - # Router is replicated: every rank computes the full top-k selection. topk_weights, topk_ids = self.gate(flat) if self._use_marlin: # Marlin's GEMM takes fp32 combine weights — pass them BEFORE the bf16 @@ -419,8 +279,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: else: topk_weights = topk_weights.to(flat.dtype) if self.packed_experts: - # Packed experts: bypass the shared bf16 ``_dispatch`` and run the - # W4A16 in-kernel dequant GEMM directly (handles tp=1 and tp>1). routed = self._dispatch_packed_experts(flat, topk_weights, topk_ids) elif self.tp_size == 1: routed = _dispatch( @@ -433,8 +291,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: ) else: routed = self._dispatch_tp(flat, topk_weights, topk_ids) - # Shared expert is a ParallelGatedMLP on the same comm group: at tp>1 it - # holds its own intermediate stripe and all-reduces inside its down_proj. shared = self.shared_expert(flat) return (routed + shared).view(input_shape) @@ -444,17 +300,6 @@ def _dispatch_packed_experts( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor: - """Packed (W4A16) routed dispatch — the memory-lean packed-expert path. - - Feeds the packed int32 weights + bf16 group scales to ``fused_experts``, - which launches ``fused_moe_kernel_w4a16`` (dequant in registers). The TP - story is identical to :meth:`_dispatch_tp`: at tp=1 the kernel sum-reduces - the top-k dim itself; at tp>1 we keep the per-slot partials - (``reduce_results=False``), all-reduce the intermediate-parallel partials, - then fold the top-k dim. Combine weights already carry - ``routed_scaling_factor`` (folded by :class:`KimiMoEGate`), so the - sum-reduce passes ``routed_scaling_factor=1.0``. - """ from mstar.utils.fused_moe import fused_experts, moe_sum_reduce_triton reduce = self.tp_size == 1 @@ -478,16 +323,7 @@ def _dispatch_packed_experts( return output def process_weights_after_loading(self, device) -> None: - """Resolve the routed-expert kernel backend and, for Marlin, repack the - loaded packed experts into Marlin layout (freeing the source packed params). - - Invoked by the generic post-load walker - (:func:`mstar.model.components.quantization.process_weights_after_loading`) - on a real device — a no-op unless the experts are packed and Marlin is both - selected (``quant_kernel != "triton"``) and eligible (sm80+, symmetric INT4, - Marlin-legal shapes). ``quant_kernel="marlin"`` raises if ineligible so an - explicit request never silently downgrades to Triton. - """ + """Resolve Triton-vs-Marlin after weights land on the real device.""" if not self.packed_experts: return from mstar.model.components.quantization import MarlinMoEMethod @@ -541,7 +377,6 @@ def process_weights_after_loading(self, device) -> None: self.experts.down_proj_scale.data, dev, ) - # Free the source packed params — Marlin holds the repacked copies now. for name in ( "gate_up_proj_packed", "gate_up_proj_scale", "down_proj_packed", "down_proj_scale", @@ -557,12 +392,6 @@ def _dispatch_marlin( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor: - """Marlin W4A16 routed dispatch. TP story is identical to - :meth:`_dispatch_packed_experts`: tp=1 sum-reduces inside the kernel; tp>1 - keeps per-slot partials (``reduce_results=False``), all-reduces the - intermediate-parallel partials, then folds the top-k dim. ``topk_weights`` - is fp32 (Marlin requirement) and already carries ``routed_scaling_factor``. - """ from mstar.utils.fused_moe import moe_sum_reduce_triton reduce = self.tp_size == 1 @@ -582,19 +411,8 @@ def _dispatch_tp( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor: - """Intermediate-sharded routed dispatch (mirrors - ``ParallelSparseMoeBlock._dispatch_tp``). - - Each rank's ``fused_experts`` produces its partial hidden contribution per - (token, top-k slot); an all-reduce sums the intermediate-dim partials - across ranks, then ``moe_sum_reduce_triton`` folds the top-k dim. The - combine weights already carry ``routed_scaling_factor`` (folded in by - :class:`KimiMoEGate`), so the sum-reduce passes ``routed_scaling_factor=1.0``. - """ from mstar.utils.fused_moe import fused_experts, moe_sum_reduce_triton - # (tokens, top_k, hidden) partials — reduce_results=False keeps the - # per-slot rows so we can all-reduce the intermediate-parallel partials. cache3 = fused_experts( flat, self.experts.gate_up_proj, diff --git a/mstar/model/kimi_k2_7/components/rope.py b/mstar/model/kimi_k2_7/components/rope.py index e170af20f..7338a59fa 100644 --- a/mstar/model/kimi_k2_7/components/rope.py +++ b/mstar/model/kimi_k2_7/components/rope.py @@ -1,18 +1,4 @@ -"""deepseek_yarn RoPE for Kimi-K2.7 / DeepSeek-V3 MLA. - -MLA rotates only the decoupled ``qk_rope_head_dim`` slice of q/k, with YARN -(NTK-by-parts) frequency scaling and an ``mscale`` amplitude on cos/sin. mstar's -``cache_manager.apply_rope`` (FlashInfer) does not implement YARN, so this is a -standalone rotary module the MLA attention applies itself. - -Style is **interleaved / GPT-J** (``is_neox_style=False`` in DeepSeek): cos/sin -are ``repeat_interleave(2)`` and adjacent even/odd pairs are rotated. Mirrors -vLLM ``DeepseekScalingRotaryEmbedding`` and the YARN helpers in ``common.py``. - -Two ``mscale`` values (both use the 2-arg ``yarn_get_mscale``): - - **amplitude** on cos/sin (here): ``get_mscale(f, mscale) / get_mscale(f, mscale_all_dim) * attn_factor``. - - **softmax-scale boost** (in the attention, not here): ``get_mscale(f, mscale_all_dim) ** 2``. -""" +"""DeepSeek YARN RoPE for Kimi-K2.7 MLA.""" from __future__ import annotations import math @@ -72,22 +58,14 @@ def __init__( super().__init__() self.rotary_dim = rotary_dim - # ``inv_freq`` is NOT a registered buffer. Buffers computed in ``__init__`` - # do not survive the production ``meta`` build -> ``to_empty(device)`` -> - # ``load_weights`` path: ``to_empty`` allocates uninitialized memory and - # never re-runs ``__init__``, and ``inv_freq`` is not in the checkpoint - # (it's derived, skipped by the loader) — so a buffer would be left as - # garbage after loading, silently corrupting YARN RoPE. Instead keep the - # scalar recipe and compute ``inv_freq`` lazily in fp32 on the target - # device (also keeps it fp32 under a bf16 model, matching DeepSeek, rather - # than being downcast by ``model.to(bf16)``). + # Do not register inv_freq: meta->to_empty leaves derived buffers + # uninitialized, and model.to(bf16) would downcast it. Recompute fp32 lazily. self._inv_freq_args = ( rotary_dim, base, factor, original_max_position_embeddings, beta_fast, beta_slow, extrapolation_factor, ) self._inv_freq_cache: torch.Tensor | None = None - # cos/sin amplitude (deepseek_scaling_rope.py:56-60). self.mscale = float( yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim) diff --git a/mstar/model/kimi_k2_7/config.py b/mstar/model/kimi_k2_7/config.py index 323a5ec12..9c0afd9b2 100644 --- a/mstar/model/kimi_k2_7/config.py +++ b/mstar/model/kimi_k2_7/config.py @@ -1,16 +1,4 @@ -"""Configuration dataclass for Kimi-K2.7 (text backbone). - -Kimi-K2.7's text architecture *is* DeepSeek-V3 — vLLM serves it as -``DeepseekV3ForCausalLM`` (``model_type: "kimi_k2"`` -> ``DeepseekV3Config``), so -this dataclass carries the full DeepSeek-V3 field set: MLA latent dims, sigmoid- -routed MoE grouping, and ``deepseek_yarn`` RoPE. - -The full-size defaults are the real ``moonshotai/Kimi-K2.7-Code`` values. That -repo is the multimodal ``KimiK25ForConditionalGeneration``; the text dims here -live NESTED under ``config.json``'s ``text_config``, and its ``quantization_config`` -is nested there too (see :meth:`k27_code` and ``kimi_model.py``). The modular tests -build from :meth:`reduced`, a tiny self-consistent config. -""" +"""Kimi-K2.7 text config, using the DeepSeek-V3 architecture fields.""" from __future__ import annotations from dataclasses import dataclass, field @@ -20,70 +8,42 @@ @dataclass class KimiK2Config: - # -- Core transformer dims -------------------------------------------- vocab_size: int = 163840 hidden_size: int = 7168 - intermediate_size: int = 18432 # dense-FFN size (first_k_dense_replace layers) + intermediate_size: int = 18432 num_hidden_layers: int = 61 num_attention_heads: int = 64 num_key_value_heads: int = 64 # MLA has no separate KV heads; kept for HF parity - rms_norm_eps: float = 1e-5 # from config.json (Kimi uses 1e-5, not DeepSeek-V3's 1e-6) + rms_norm_eps: float = 1e-5 max_position_embeddings: int = 262144 tie_word_embeddings: bool = False hidden_act: str = "silu" - # -- MLA (Multi-head Latent Attention) latent dims -------------------- - # Query is compressed to ``q_lora_rank`` then projected up to - # ``num_attention_heads * qk_head_dim``; K/V share a ``kv_lora_rank`` latent - # plus a decoupled ``qk_rope_head_dim`` RoPE slice. Per-head query/key dim is - # ``qk_nope_head_dim + qk_rope_head_dim``; value head dim differs. q_lora_rank: int = 1536 kv_lora_rank: int = 512 qk_nope_head_dim: int = 128 qk_rope_head_dim: int = 64 v_head_dim: int = 128 - # -- MLA weight absorption (DEFAULT) ----------------------------------- - # ``True`` (default) => weight-absorbed MLA: ``kv_b_proj``'s up-projection is - # folded into the Q path (``W_UK``) and the O path (``W_UV``) at load (plus the - # ``fused_qkv_a_proj`` down-proj fusion), attention runs as MQA over the - # COMPRESSED latent via the ``mla_absorb`` cache backend, and the KV cache - # stores only the ``kv_lora_rank + qk_rope_head_dim`` latent (1 KV head) — a - # ~57x per-token cache shrink, numerically identical up to fp rounding. - # ``False`` => naive/materialized MLA (latent projected up to full per-head - # K/V, padded to ``padded_head_dim``, MHA cache): the M4-golden parity - # reference / opt-out fallback. - # - # PERF CAVEAT: the absorbed backend currently runs on a torch SDPA-over-latent - # path — correct + memory-lean but EAGER-ONLY (no CUDA-graph capture) and slow - # on the real 1T. The FlashInfer MLA kernel + CUDA-graph capture (production - # throughput) is a follow-up; until it lands, real large-scale serving should - # set ``mla_absorb=False`` (naive) or accept eager execution. See - # ``components/attention.py``, ``kimi_model.py::get_kv_cache_config``, - # ``engine/cache_manager.py::MlaAbsorbCacheManager``. + # Default absorbed MLA stores one compressed latent KV head. The naive path is + # kept as the reduced-test parity fallback. mla_absorb: bool = True - # -- Fine-grained MoE (sigmoid router, group-limited top-k, noaux_tc) -- - n_routed_experts: int = 384 # from config.json + n_routed_experts: int = 384 n_shared_experts: int = 1 - num_experts_per_tok: int = 8 # top-k + num_experts_per_tok: int = 8 moe_intermediate_size: int = 2048 - n_group: int = 1 # from config.json - topk_group: int = 1 # from config.json (groups kept by group-limited routing) - routed_scaling_factor: float = 2.827 # from config.json - scoring_func: str = "sigmoid" # DeepSeek-V3/Kimi: sigmoid (not softmax) - topk_method: str = "noaux_tc" # per-expert e_score_correction_bias + n_group: int = 1 + topk_group: int = 1 + routed_scaling_factor: float = 2.827 + scoring_func: str = "sigmoid" + topk_method: str = "noaux_tc" norm_topk_prob: bool = True - first_k_dense_replace: int = 1 # first N layers are dense, rest are MoE + first_k_dense_replace: int = 1 moe_layer_freq: int = 1 - # -- deepseek_yarn RoPE ------------------------------------------------ - rope_theta: float = 50000.0 # from config.json + rope_theta: float = 50000.0 rope_scaling: dict = field(default_factory=lambda: { - # from config.json (HF key is "type": "yarn"; mstar's internal id for the - # DeepSeek/Kimi variant is "deepseek_yarn"). factor=64 yields the 262144 - # context (4096 * 64). K2.7-Code keeps beta_fast=32; mscale == - # mscale_all_dim == 1.0. "rope_type": "deepseek_yarn", "factor": 64.0, "original_max_position_embeddings": 4096, @@ -93,76 +53,35 @@ class KimiK2Config: "mscale_all_dim": 1.0, }) - # -- Special tokens / generation defaults ----------------------------- - bos_token_id: int = 163584 # from config.json - eos_token_id: int = 163586 # from config.json - pad_token_id: int = 163839 # from config.json + bos_token_id: int = 163584 + eos_token_id: int = 163586 + pad_token_id: int = 163839 temperature: float = 1.0 top_p: float = 1.0 ignore_eos: bool = False - # -- MTP (multi-token prediction) — deferred, declared for completeness - num_nextn_predict_layers: int = 0 - # -- Quantization (compressed-tensors INT4/fp8) ----------------------- - # ``None`` => native-bf16 checkpoint. When set, the weight loader dequantizes - # the checkpoint stream on load (:mod:`mstar.model.kimi_k2_7.quantization`) - # before the bf16 remap + stacked rules. Populated from the real checkpoint's - # ``config.json`` ``quantization_config`` (``kimi_model.py``) or set directly - # for the reduced/synthetic tests (:meth:`reduced_quantized`). quantization_config: CompressedTensorsQuantConfig | None = None - # -- Quantization: memory-lean packed experts (in-kernel dequant) ---------- - # ``False`` => quantized routed experts are dequantized to bf16 on load and fed - # to the bf16 fused-expert GEMM. ``True`` (only meaningful when - # ``quantization_config`` is set) => the routed experts stay PACKED int32 in - # VRAM and the W4A16 ``fused_moe_kernel_w4a16`` dequantizes each tile in - # registers. MLA / dense-FFN / shared-expert weights are always dequantized on - # load. This is the only path that fits the real 1T checkpoint. See - # ``components/moe.py`` / ``weight_loader.py``. + # Keeps routed experts packed; non-expert quantized weights still dequantize + # on load. moe_in_kernel_dequant: bool = False - # -- Quantization: routed-expert W4A16 kernel backend ---------------------- - # Chooses the kernel for the PACKED routed experts (only meaningful when - # ``moe_in_kernel_dequant`` is set — Marlin layers on top of the Hook B packed - # params). Values: - # "auto" => Marlin on sm80+ (Ampere/Hopper) when the build succeeds and the - # shapes/group_size are Marlin-legal, else the Triton - # ``fused_moe_kernel_w4a16`` fallback. This is the production default. - # "marlin" => force Marlin; raise if ineligible (must not silently downgrade). - # "triton" => force the Triton in-kernel dequant path (the pre-Marlin behavior). - # The final resolution happens post-load in - # ``KimiSparseMoeBlock.process_weights_after_loading`` (a real device is needed to - # probe capability + build the kernel); ``__init__`` runs on ``meta``. + # "auto" probes Marlin post-load on the real device; "marlin" must not silently + # downgrade, and "triton" keeps the packed Triton path. quant_kernel: str = "auto" - # -- Serving: CUDA-graph prefill capture grid (optional overrides) ------ - # ``None`` => ``KimiLLMSubmodule`` uses its full-size class-default grid. - # ``reduced()`` sets a tiny grid so the synthetic bring-up serve captures a - # single short-prompt graph instead of the full 6x5 compiled grid. prefill_token_buckets: list[int] | None = None prefill_capture_batch_sizes: list[int] | None = None - # --------------------------------------------------------------------- - # Derived dims (read by get_kv_cache_config / attention) - # --------------------------------------------------------------------- @property def qk_head_dim(self) -> int: - """Per-head query/key dim: nope + decoupled-rope slice (e.g. 128+64=192).""" return self.qk_nope_head_dim + self.qk_rope_head_dim @property def padded_head_dim(self) -> int: - """Head dim the naive-MLA q/k/v are zero-padded to for the paged cache. - - FlashInfer's SM90 (Hopper) prefill kernel ``static_assert``s - ``head_dim_vo ∈ {64, 128, 256}``, so it will not JIT-build for the real - ``qk_head_dim=192`` or the reduced ``qk_head_dim=24``. We pad q/k (from - ``qk_head_dim``) and v (from ``v_head_dim``) up to the smallest supported - dim ``>= qk_head_dim``, run the paged attention there, and slice the output - back to ``v_head_dim`` — compensating the softmax scale (see - ``KimiMLAAttention.softmax_scale_boost``). Real Kimi 192 -> 256; reduced 24 -> 64. - """ + """Naive-MLA q/k/v pad target; FlashInfer paged kernels require 64/128/256.""" for supported in (64, 128, 256): if supported >= self.qk_head_dim: return supported @@ -177,17 +96,6 @@ def num_dense_layers(self) -> int: @classmethod def reduced(cls) -> "KimiK2Config": - """A tiny, self-consistent config for CPU/dummy-mode modular tests and - reduced-config golden runs. Keeps the *shape* of Kimi (MLA split heads, - grouped MoE, one dense layer) while being small enough to run without - the 1T checkpoint. - - NOTE ``mla_absorb=False``: this fixture pins the NAIVE MLA path, the - M4-golden parity reference that the bulk of the reduced test suite - validates. The absorbed path (the production default) is exercised by the - dedicated ``test_kimi_mla_absorb*`` tests, which flip ``mla_absorb=True`` - on a reduced() instance explicitly. - """ return cls( mla_absorb=False, vocab_size=256, @@ -209,10 +117,6 @@ def reduced(cls) -> "KimiK2Config": n_group=1, topk_group=1, first_k_dense_replace=1, - # Tiny CUDA-graph prefill capture grid for the synthetic bring-up serve: - # one short-prompt bucket at batch size 1 (the full 6x5 grid is slow and - # its larger buckets exceed this 512-token model). Serve/CUDA-graph path - # only — the golden tests call forward() directly and are unaffected. prefill_token_buckets=[64], prefill_capture_batch_sizes=[1], ) @@ -224,15 +128,6 @@ def reduced_quantized( group_size: int = 32, symmetric: bool = True, ) -> "KimiK2Config": - """:meth:`reduced` plus a compressed-tensors quant config, to exercise the - dequant-on-load path on a synthetic quantized checkpoint. - - The reduced dims (``hidden_size=128``, ``moe_intermediate_size=64``, - ``intermediate_size=256`` …) are all divisible by the default - ``group_size=32`` and by ``pack_factor=8``, so the FFN / expert / MLA - weights whose input dim divides ``group_size`` can be quantized while the - rest stay bf16 — the mixed checkpoint the streaming parser handles. - """ cfg = cls.reduced() cfg.quantization_config = CompressedTensorsQuantConfig( num_bits=num_bits, group_size=group_size, symmetric=symmetric, @@ -246,15 +141,6 @@ def reduced_quantized_inkernel( group_size: int = 32, symmetric: bool = True, ) -> "KimiK2Config": - """:meth:`reduced_quantized` plus ``moe_in_kernel_dequant=True`` — packed - routed experts + in-kernel INT4 dequant on a synthetic quantized checkpoint. - - The reduced dims (``hidden_size=128``, ``moe_intermediate_size=64``) satisfy - the packed-expert divisibility asserts (``% pack_factor`` and ``% - group_size``) at tp=1 (``shard_inter=64``) and tp=2 (``shard_inter=32``); - tp=4 (``shard_inter=16``) fails ``% group_size`` (32), so pin packed-expert - goldens to tp<=2. - """ cfg = cls.reduced_quantized( num_bits=num_bits, group_size=group_size, symmetric=symmetric, ) @@ -268,18 +154,6 @@ def reduced_marlin( group_size: int = 32, symmetric: bool = True, ) -> "KimiK2Config": - """:meth:`reduced_quantized_inkernel` with Marlin-legal shapes + the Marlin - routed-expert backend forced on. - - Marlin's GEMM imposes ``n % 64 == 0`` and ``k % 128 == 0`` on each expert - matmul, which the default reduced dims (``hidden_size=128``, - ``moe_intermediate_size=64``) do NOT satisfy for the down projection - (``k == shard_inter``). This variant bumps ``hidden_size=256`` and - ``moe_intermediate_size=256`` so both expert GEMMs are Marlin-legal at - tp<=2 (tp=1 ``shard_inter=256``, tp=2 ``shard_inter=128`` — both ``% 128``); - tp=4 (``shard_inter=64``) fails ``k % 128``, so pin Marlin goldens to tp<=2. - ``group_size=32`` and ``pack_factor=8`` still divide both axes. - """ cfg = cls.reduced_quantized_inkernel( num_bits=num_bits, group_size=group_size, symmetric=symmetric, ) @@ -291,19 +165,6 @@ def reduced_marlin( @classmethod def k27_code(cls) -> "KimiK2Config": - """Full-size ``moonshotai/Kimi-K2.7-Code`` text-only serve config. - - Full 1T dims plus ``moe_in_kernel_dequant=True``: Kimi-K2.7-Code is a ~1T - INT4 ``pack-quantized`` checkpoint (num_bits=4, group_size=32, symmetric, - routed experts only), so the routed experts are served packed and - dequantized in-kernel — dequantizing them to bf16 would need ~2 TB of VRAM. - The ``quantization_config`` is nested under ``text_config`` and auto-read at - load by ``kimi_model.py::_maybe_apply_checkpoint_quant_config``; MLA / - dense-FFN / shared-expert / lm_head / vision weights stay bf16, matching the - checkpoint ``ignore`` list. - - Keeps the default ``beta_fast=32.0``. - """ cfg = cls() cfg.moe_in_kernel_dequant = True return cfg diff --git a/mstar/model/kimi_k2_7/kimi_model.py b/mstar/model/kimi_k2_7/kimi_model.py index 7c00fe605..9d6da4883 100644 --- a/mstar/model/kimi_k2_7/kimi_model.py +++ b/mstar/model/kimi_k2_7/kimi_model.py @@ -1,18 +1,4 @@ -"""KimiK2Model: M* Model contract for Kimi-K2.7 (text backbone). - -Kimi-K2.7's text path is DeepSeek-V3 (``model_type: "kimi_k2"`` -> -``DeepseekV3ForCausalLM``). This declares the full serving plumbing — the graph -(``prefill`` + ``decode`` Loop), the single ``KV_CACHE`` LLM node, the KV-cache -dims, and the prefill->decode->done state machine — and builds the LLM submodule -in ``get_submodule``. When ``get_submodule`` returns ``None`` (dummy mode), -``pytest test/modular/`` exercises the graph/walk/engine-routing machinery in -isolation, as ``docs/adding_models.rst`` prescribes. - -Structurally this mirrors Orpheus's LLM partition (the smallest complete LLM in -the tree) minus the async SNAC partition: Kimi text-only is a single ``default`` -partition, so it inherits ``Model.get_partitions`` / ``get_partition_topology`` -and only implements the abstract surface. -""" +"""M* model wrapper for the Kimi-K2.7 text backbone.""" from __future__ import annotations import logging @@ -40,7 +26,6 @@ def _resolve_local_hf_snapshot(repo_id: str, cache_dir: str | None = None) -> str: - """Resolve an HF repo id to a local snapshot dir (mirrors OrpheusModel).""" from pathlib import Path from huggingface_hub import snapshot_download @@ -56,8 +41,6 @@ def _resolve_local_hf_snapshot(repo_id: str, cache_dir: str | None = None) -> st class KimiK2Model(Model): - """Kimi-K2.7 text backbone (DeepSeek-V3 architecture).""" - def __init__( self, model_path_hf: str, @@ -65,40 +48,20 @@ def __init__( **kwargs, ): self.cache_dir = cache_dir - # ``model_kwargs`` from the serving YAML arrive here as ``**kwargs`` (see - # api_server/entrypoint.py). They let a config redirect this model at a - # local (reduced/synthetic) checkpoint without touching the shared - # registry, so a runnable text serve is possible before the 1T weights - # exist. All three are optional and default to the full-size behaviour. - # * ``checkpoint_path`` — local HF-format dir/file to load instead of - # the ``HF_MODELS`` repo id (used as-is by ``_resolve_checkpoint``). - # * ``config_variant`` — ``"reduced"`` selects ``KimiK2Config.reduced()`` - # (tiny, GPU-runnable shape); anything else keeps the 1T config. - # * ``tokenizer_mode`` — ``"byte"`` swaps the HF tokenizer for a trivial - # UTF-8 byte identity tokenizer, the pragmatic fit for the reduced - # ``vocab_size=256`` model (the real Kimi tokenizer emits ids ≫ 256). checkpoint_path = kwargs.get("checkpoint_path") self.model_path_hf = checkpoint_path or model_path_hf self._config_variant = kwargs.get("config_variant", "full") if self._config_variant == "reduced": self.config = KimiK2Config.reduced() elif self._config_variant == "reduced_quantized": - # Reduced shape + a quant config, to exercise dequant-on-load. self.config = KimiK2Config.reduced_quantized() elif self._config_variant == "reduced_quantized_inkernel": - # Reduced shape + quant config + packed experts (in-kernel W4A16 dequant). - # int32 packed params are auto-exempt from the whole-model ``.to(bf16)`` - # cast below (PyTorch ``.to(dtype)`` only casts float/complex), no hook. self.config = KimiK2Config.reduced_quantized_inkernel() elif self._config_variant == "k27_code": - # Full-size Kimi-K2.7-Code text-only serve config (see KimiK2Config.k27_code). self.config = KimiK2Config.k27_code() else: self.config = KimiK2Config() self._tokenizer_mode = kwargs.get("tokenizer_mode", "hf") - # Tokenizer is loaded lazily: the modular (dummy-mode) tests build the - # model via ``object.__new__`` and never call ``__init__``, so we avoid - # forcing a network/tokenizer dependency into the scaffold path. self._tokenizer = None self._submodule_cache: dict[str, NodeSubmodule | None] = {} @@ -114,21 +77,8 @@ def tokenizer(self): ) return self._tokenizer - # ------------------------------------------------------------------- - # Model ABC: KV cache config - # ------------------------------------------------------------------- - def get_kv_cache_config(self) -> list[KVCacheConfig]: if self.config.mla_absorb: - # Weight-absorbed MLA (the default): attention is MQA over the - # COMPRESSED latent (kv_b_proj folded into Q/O + the fused_qkv_a_proj - # down-proj), so the paged cache stores a single KV "head" of width - # ``kv_lora_rank + qk_rope_head_dim`` ([kv_c | k_pe]) per token — a ~57x - # shrink vs the naive padded MHA cache (real 2*64*256=32768 -> 512+64=576; - # reduced 2*4*64=512 -> 40). Served by ``MlaAbsorbCacheManager`` over the - # 4D latent cache. ``softmax_scale`` is DeepSeek's intended MLA scale - # ``qk_head_dim**-0.5 * mscale**2``: the absorbed forward folds nothing - # into q (unlike naive), so the backend applies this scale directly. from mstar.model.kimi_k2_7.components.rope import yarn_get_mscale rope = self.config.rope_scaling mscale = yarn_get_mscale(rope["factor"], rope.get("mscale_all_dim", 0.0)) @@ -141,20 +91,8 @@ def get_kv_cache_config(self) -> list[KVCacheConfig]: num_qo_heads=self.config.num_attention_heads, attention_backend="mla_absorb", softmax_scale=softmax_scale, - # The ckv/kpe split for the FlashInfer MLA kernel fast path - # (head_dim = ckv + kpe = kv_lora_rank + qk_rope_head_dim). mla_ckv_dim=self.config.kv_lora_rank, )] - # Naive/materialized MLA (the first-pass port, per CLAUDE.md): the latent - # is projected up to full per-head K/V and broadcast to every query head, - # so from the paged cache's ``[tokens, heads, head_dim]`` point of view - # there are ``num_attention_heads`` KV heads. K/V are stored at - # ``padded_head_dim`` — the naive path zero-pads q/k (from ``qk_head_dim``, - # e.g. 192) and v (from ``v_head_dim``) up to the smallest FlashInfer-SM90 - # supported head_dim >= qk_head_dim (256 real, 64 reduced), because the - # Hopper prefill kernel static_asserts head_dim_vo in {64,128,256}. The - # attention output is sliced back to ``v_head_dim`` in the submodule. This - # trades cache size for not needing a weight-absorb path in the engine. return [KVCacheConfig( num_layers=self.config.num_hidden_layers, num_kv_heads=self.config.num_attention_heads, @@ -163,21 +101,10 @@ def get_kv_cache_config(self) -> list[KVCacheConfig]: num_qo_heads=self.config.num_attention_heads, )] - # ------------------------------------------------------------------- - # Model ABC: node engine types - # ------------------------------------------------------------------- - def get_node_engine_types(self) -> dict[str, EngineType]: return {LLM_NODE: EngineType.KV_CACHE} - # ------------------------------------------------------------------- - # Model ABC: graph walk definitions - # ------------------------------------------------------------------- - def get_graph_walk_graphs(self) -> dict[str, GraphSection]: - # prefill: embed the prompt, fill the KV cache, sample + emit the first - # token. ``persist=True`` keeps that token at the conductor so the decode - # walk can pick it up as its first ``text_inputs``. prefill = GraphNode( name=LLM_NODE, input_names=["text_inputs"], @@ -192,10 +119,6 @@ def get_graph_walk_graphs(self) -> dict[str, GraphSection]: ], ) - # decode: autoregressive Loop. Each step emits the new token to the - # client and feeds it back as the next step's ``text_inputs``. The Loop - # stops via the submodule's ``check_stop`` (EOS / max tokens); ``max_iters`` - # is the hard cap. decode = Loop( name=DECODE_LOOP, section=GraphNode( @@ -220,10 +143,6 @@ def get_graph_walk_graphs(self) -> dict[str, GraphSection]: return dict(prefill=prefill, decode=decode) - # ------------------------------------------------------------------- - # Model ABC: forward pass args (single "default" partition) - # ------------------------------------------------------------------- - def get_initial_forward_pass_args( self, partition_name: str, @@ -258,14 +177,6 @@ def get_partition_forward_pass_args( persist_signals: dict[str, list[TensorPointerInfo]], incoming_connections: list[StreamingConnectionState] | None = None, ) -> ForwardPassArgs: - """Drive the prefill → decode → done state machine. - - Called by the conductor after each completed walk. Prefill transitions to - the decode walk (feeding the persisted first token as ``text_inputs``); - once the decode walk completes (its Loop stopped via ``check_stop``), the - request is done. The per-token decode iteration is driven inside the - graph Loop, not by repeated calls here. - """ metadata = partition_metadata request_done = False @@ -296,10 +207,6 @@ def get_partition_forward_pass_args( step_metadata={"is_prefill": metadata.is_prefill}, ) - # ------------------------------------------------------------------- - # Model ABC: prompt processing - # ------------------------------------------------------------------- - def process_prompt( self, prompt: str | None, @@ -308,14 +215,10 @@ def process_prompt( tensors: NameToTensorList | None = None, **kwargs, ) -> NameToTensorList: - # Text-only for M0; raw multimodal tensors (MoonViT) are a later milestone. if prompt is None: return {} if self._tokenizer_mode == "byte": - # Trivial UTF-8 byte identity tokenizer for the reduced vocab_size=256 - # serve: each prompt byte is already a valid token id in [0, 256), so - # no HF tokenizer / network dependency is needed. Clamped defensively - # in case a smaller reduced vocab is ever used. + # Reduced serve maps UTF-8 bytes directly to token ids, avoiding HF IO. vocab = self.config.vocab_size byte_ids = [min(b, vocab - 1) for b in prompt.encode("utf-8")] or [0] input_ids = torch.tensor(byte_ids, dtype=torch.long) @@ -336,10 +239,6 @@ def get_sampling_config( ignore_eos=model_kwargs.get("ignore_eos", self.config.ignore_eos), ) - # ------------------------------------------------------------------- - # Model ABC: postprocess - # ------------------------------------------------------------------- - def postprocess( self, output: torch.Tensor, @@ -349,32 +248,17 @@ def postprocess( if modality == "text": token_ids = output.tolist() if output.numel() else [] if self._tokenizer_mode == "byte": - # Inverse of the byte identity tokenizer: reduced-vocab ids map - # straight back to raw bytes. The synthetic model emits arbitrary - # ids in [0, 256), so the bytes are not guaranteed valid UTF-8 — - # decode leniently (the point is to prove tokens stream, not to - # produce meaningful text on random weights). + # Synthetic reduced models emit arbitrary byte ids; return raw bytes. return bytes((t & 0xFF) for t in token_ids) text = self.tokenizer.decode(token_ids, skip_special_tokens=True) return text.encode("utf-8") raise ValueError(f"Unsupported modality for Kimi-K2.7: {modality!r}") - # ------------------------------------------------------------------- - # Model ABC: sharding - # ------------------------------------------------------------------- - def get_default_sharding_config(self): from mstar.distributed.base import ShardingConfig - # Kimi is a 1T model — real serving is TP8 / multi-node. The LLM node is - # the tensor-parallel node; the per-node degree comes from the config - # YAML's ``node_groups``, not from the model code. return ShardingConfig(groups=[], tp_enabled_nodes={LLM_NODE}, shard_dim={}) - # ------------------------------------------------------------------- - # Model ABC: submodule loading - # ------------------------------------------------------------------- - def get_submodule( self, node_name: str, @@ -403,25 +287,14 @@ def _create_submodule( source = self._resolve_checkpoint() if source is None: - # Dummy mode: no checkpoint resolvable (e.g. the modular graph tests - # build the model via object.__new__ with no model_path_hf). Returning - # None lets pytest test/modular/ validate the graph/walks/engine-routing - # without a GPU or weights, per docs/adding_models.rst. logger.info( "KimiK2Model: no checkpoint resolved for node %r — dummy mode (None).", node_name, ) return None - # If the checkpoint declares a compressed-tensors ``quantization_config``, - # route loading through the dequant-on-load parser (weight_loader). An - # explicit config (e.g. reduced_quantized) is respected and not clobbered. self._maybe_apply_checkpoint_quant_config(source) - # Real build, mirroring OrpheusModel._create_llm_submodule: construct on the - # meta device (no allocation), cast to the target dtype on meta (so to_empty - # allocates directly in bf16, not fp32-then-downcast), materialise storage, - # then run the HF loader (remap + fused-expert stacked rules). from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM from mstar.model.kimi_k2_7.submodules import KimiLLMSubmodule from mstar.model.loader import load_weights @@ -432,13 +305,6 @@ def _create_submodule( language_model = language_model.to(autocast_dtype) language_model.to_empty(device=device) load_weights(language_model, source, device=device) - # Post-load pass: let any submodule finalize its weights on the real device - # now that the checkpoint is resident. The routed-expert MoE block repacks - # its packed experts into the Marlin (or Triton) kernel layout + allocates a - # workspace; under mla_absorb, KimiMLAAttention builds the absorbed w_kc/w_vc - # projections + the fused_qkv_a_proj weight. Generic + idempotent walker - # (calls the hook on every module that exposes one); no-op for a plain bf16 - # module without the hook. from mstar.model.components.quantization import process_weights_after_loading process_weights_after_loading(language_model, torch.device(device)) @@ -448,12 +314,6 @@ def _create_submodule( return KimiLLMSubmodule(language_model=language_model, config=self.config) def _resolve_checkpoint(self) -> str | None: - """Resolve the HF checkpoint source, or None for dummy mode. - - A local directory / file (e.g. a reduced synthetic checkpoint) is used - as-is; otherwise the HF repo id is snapshot-downloaded. Returns None when - no ``model_path_hf`` is set (dummy-mode graph tests). - """ from pathlib import Path path = getattr(self, "model_path_hf", None) @@ -464,20 +324,6 @@ def _resolve_checkpoint(self) -> str | None: return _resolve_local_hf_snapshot(path, cache_dir=getattr(self, "cache_dir", None)) def _maybe_apply_checkpoint_quant_config(self, source: str) -> None: - """Populate ``self.config.quantization_config`` from ``config.json``. - - Reads the checkpoint's ``config.json`` ``quantization_config`` block (a - compressed-tensors INT4/fp8 checkpoint carries one) and stores the parsed - :class:`CompressedTensorsQuantConfig` on the model config so the weight - loader takes the dequant-on-load path. A config set explicitly - (e.g. ``reduced_quantized()``) wins and is left untouched; a plain bf16 - checkpoint (no block, unreadable, or single-file source) is a no-op. - - The real multimodal ``Kimi-K2.7-Code`` repo nests the block under - ``text_config`` (the top-level ``quantization_config`` is null), so this - reads a top-level block if present, otherwise ``text_config``'s. - ``from_hf_config_dict`` parses the same block shape either way. - """ import json from pathlib import Path @@ -494,8 +340,6 @@ def _maybe_apply_checkpoint_quant_config(self, source: str) -> None: except (OSError, ValueError) as e: # unreadable / malformed — stay bf16 logger.warning("KimiK2Model: could not read %s: %s", config_json, e) return - # A top-level ``quantization_config`` if present, else the one nested under - # ``text_config`` (the multimodal K2.7-Code layout — top-level is null). quant_raw = raw.get("quantization_config") or ( raw.get("text_config") or {} ).get("quantization_config") diff --git a/mstar/model/kimi_k2_7/quantization.py b/mstar/model/kimi_k2_7/quantization.py index d0ff5085e..7f2ffef22 100644 --- a/mstar/model/kimi_k2_7/quantization.py +++ b/mstar/model/kimi_k2_7/quantization.py @@ -1,22 +1,7 @@ -"""Compressed-tensors INT4 (W4A16) quantization for Kimi-K2.7 weights. +"""Compressed-tensors INT4 W4A16 helpers for Kimi-K2.7 weights. -On-disk format (compressed-tensors ``pack-quantized``), per quantized Linear -weight of logical shape ``(out, in)``: - - * ``.weight_packed`` — int32 ``(out, in // pack_factor)``; ``pack_factor = - 32 // num_bits`` (8 for INT4) values packed low-order-first along the input axis. - * ``.weight_scale`` — bf16 ``(out, in // group_size)``, one scale per (row, group). - * ``.weight_zero_point`` — asymmetric only. - * ``.weight_shape`` — original ``(out, in)``; optional, used only to validate. - -Symmetric INT4 is stored offset-binary: the packed nibble is ``(signed value + 8)``, -so dequant subtracts 8 (matches vLLM's ``uint4b8``); asymmetric subtracts the zero -point. To flip a checkpoint to plain two's-complement, change the ``bias`` line in -:func:`dequantize_weight`. Layout authority: vLLM -``compressed_tensors/schemes/compressed_tensors_wNa16.py``. - -``dequant_compressed_tensors_stream`` dequantizes a checkpoint stream to bf16 on -load; ``keep_packed`` leaves selected weights packed for in-kernel dequant. +Packed values are stored low-order-first in int32 containers. Symmetric INT4 is +offset-binary, so dequant subtracts 8 to match vLLM's ``uint4b8`` layout. """ from __future__ import annotations @@ -25,7 +10,6 @@ import torch -# Suffixes a compressed-tensors checkpoint attaches to each quantized tensor. _PACKED = ".weight_packed" _SCALE = ".weight_scale" _ZERO_POINT = ".weight_zero_point" @@ -35,14 +19,10 @@ @dataclass(frozen=True) class CompressedTensorsQuantConfig: - """The subset of a compressed-tensors ``quantization_config`` this port reads. - - Kimi-K2.7 ships a single ``config_groups`` entry (``weights`` only, W4A16), so - the whole checkpoint shares one ``num_bits`` / ``group_size`` / ``symmetric``. - """ + """Subset of the compressed-tensors config used by Kimi-K2.7.""" num_bits: int = 4 - group_size: int = 32 # -1 => channelwise (one group spans the full input dim) + group_size: int = 32 symmetric: bool = True strategy: str = "group" # "group" | "channel" quant_format: str = "pack-quantized" @@ -51,19 +31,12 @@ class CompressedTensorsQuantConfig: @property def pack_factor(self) -> int: - """Number of ``num_bits`` values packed into one int32.""" return 32 // self.num_bits @classmethod def from_hf_config_dict( cls, quant: dict | None ) -> "CompressedTensorsQuantConfig | None": - """Build from a checkpoint ``config.json``'s ``quantization_config`` block. - - Returns ``None`` when there is no quantization block (a plain bf16 - checkpoint). Reads the first ``config_groups`` entry's ``weights`` spec — - Kimi uses exactly one group. - """ if not quant: return None groups = quant.get("config_groups") or {} @@ -85,21 +58,8 @@ def from_hf_config_dict( ignore=tuple(quant.get("ignore", []) or []), ) - -# --------------------------------------------------------------------------- -# Bit packing — exact inverses. Packing is along the last (input) axis, which is -# the input axis of a checkpoint ``(out, in)`` Linear weight (what gets quantized). -# --------------------------------------------------------------------------- - def pack_int32(values_unsigned: torch.Tensor, num_bits: int) -> torch.Tensor: - """Pack ``pack_factor`` unsigned ``num_bits`` values (last axis) into int32. - - ``values_unsigned`` holds integers in ``[0, 2**num_bits)``; the last axis must - be divisible by ``pack_factor = 32 // num_bits``. Values are combined - low-order-first (element ``j`` occupies bits ``[num_bits*j, num_bits*(j+1))``), - matching compressed-tensors. Returns int32 of shape - ``(..., last // pack_factor)``. - """ + """Pack unsigned values along the last axis into low-order-first int32.""" pack_factor = 32 // num_bits *lead, n = values_unsigned.shape if n % pack_factor != 0: @@ -107,17 +67,11 @@ def pack_int32(values_unsigned: torch.Tensor, num_bits: int) -> torch.Tensor: q = values_unsigned.to(torch.int64).reshape(*lead, n // pack_factor, pack_factor) shifts = torch.arange(pack_factor, device=q.device, dtype=torch.int64) * num_bits packed = (q << shifts).sum(dim=-1) - # Wrap the 32-bit pattern into a signed int32 container (matches on-disk dtype). return (packed & 0xFFFFFFFF).to(torch.int32) def unpack_int32(packed: torch.Tensor, num_bits: int) -> torch.Tensor: - """Inverse of :func:`pack_int32`: expand int32 to unsigned ``num_bits`` nibbles. - - Returns an int64 tensor of shape ``(..., last * pack_factor)`` with values in - ``[0, 2**num_bits)``. The int32 is read as an unsigned 32-bit pattern, so the - top nibble is recovered correctly regardless of the container's sign bit. - """ + """Inverse of :func:`pack_int32`; reads int32 as an unsigned bit pattern.""" pack_factor = 32 // num_bits mask = (1 << num_bits) - 1 p = packed.to(torch.int64) & 0xFFFFFFFF @@ -137,21 +91,7 @@ def dequantize_weight( zero_point: torch.Tensor | None = None, out_dtype: torch.dtype = torch.bfloat16, ) -> torch.Tensor: - """Dequantize one compressed-tensors weight to ``out_dtype``. - - Args: - packed: ``(out, in // pack_factor)`` int32 packed weight. - scale: ``(out, in // group_size)`` per-(row, group) scale. - num_bits: bit width (4 for Kimi INT4). - group_size: group granularity along the input axis; ``-1`` => channelwise. - symmetric: symmetric offset-binary (subtract ``2**(num_bits-1)``) vs. - asymmetric (subtract ``zero_point``). - zero_point: ``(out, in // group_size)`` per-group zero point (asymmetric). - out_dtype: result dtype (bf16 to feed the existing fused-expert GEMM). - - Returns: - ``(out, in)`` dequantized weight, in ``out_dtype``. - """ + """Dequantize one packed compressed-tensors weight.""" nibbles = unpack_int32(packed, num_bits).to(torch.float32) # (out, in) unsigned out_f, in_f = nibbles.shape gs = in_f if group_size in (-1, None) else group_size @@ -172,50 +112,23 @@ def dequantize_weight( if s.shape[-1] != in_f: # per-group -> broadcast to per-column s = s.repeat_interleave(gs, dim=-1) return (nibbles * s).to(out_dtype) - -# --------------------------------------------------------------------------- -# Streaming dequant-on-load — the generator wired into load_kimi_hf_weights. -# --------------------------------------------------------------------------- - def dequant_compressed_tensors_stream( weights: Iterable[tuple[str, torch.Tensor]], quant_config: CompressedTensorsQuantConfig, out_dtype: torch.dtype = torch.bfloat16, keep_packed: Callable[[str], bool] | None = None, ) -> Iterator[tuple[str, torch.Tensor]]: - """Wrap a checkpoint ``(name, tensor)`` stream, dequantizing on the fly. - - For every quantized tensor the checkpoint carries ``.weight_packed`` + - ``.weight_scale`` (+ ``.weight_zero_point`` for asymmetric); this - buffers those components per ```` and, once complete, yields a single - ``(.weight, bf16 tensor)`` — exactly the key a native-bf16 checkpoint - would carry — then drops the quant sub-keys. Any key that is not a - compressed-tensors component (norms, the router ``gate``, ``embed_tokens``, - ``lm_head``, or a weight the checkpoint left in bf16) passes straight through. - - Buffering is bounded to the in-flight incomplete tensors: a ```` is - emitted and freed the moment its required components have all been seen, - independent of the iterator's key order. - - ``keep_packed``: when it returns True for a ````, that base's - compressed-tensors sub-keys are passed through RAW (no buffering, no dequant) - so a downstream packed-expert loader can route them to int32 params. Kimi - passes a predicate matching the routed experts (kept packed for in-kernel - dequant) while every other quantized weight — MLA, dense FFN, shared expert — - still dequantizes here. Because kept bases never enter ``buffers``, the - end-of-stream completeness check is unaffected. - """ + """Convert quantized sub-key streams to bf16 weights unless ``keep_packed`` matches.""" buffers: dict[str, dict[str, torch.Tensor]] = {} for name, tensor in weights: suffix = next((s for s in _QUANT_SUFFIXES if name.endswith(s)), None) if suffix is None: - yield name, tensor # not a quant component — pass through untouched + yield name, tensor continue base = name[: -len(suffix)] if keep_packed is not None and keep_packed(base): - # Leave this base packed — hand the raw sub-key downstream. yield name, tensor continue diff --git a/mstar/model/kimi_k2_7/submodules.py b/mstar/model/kimi_k2_7/submodules.py index 49294b042..907b5fba3 100644 --- a/mstar/model/kimi_k2_7/submodules.py +++ b/mstar/model/kimi_k2_7/submodules.py @@ -1,21 +1,4 @@ -"""Submodules for Kimi-K2.7 (text backbone). - -The :class:`KimiLLMSubmodule` ``ARNodeSubmodule`` drives the DeepSeek-V3 text -backbone (MLA attention over the paged cache + fine-grained sigmoid-routed MoE) -through the engine's ``prepare_inputs -> preprocess -> -forward/forward_batched -> postprocess -> check_stop`` lifecycle for the -``prefill`` and ``decode`` Loop walks. - -Structurally this mirrors ``OrpheusLLMSubmodule`` (the smallest complete LLM in -the tree) with one addition: the naive MLA applies YARN RoPE itself over the -decoupled ``qk_rope`` slice, so ``preprocess`` builds per-token ``position_ids`` -(the same positions ``plan_rope`` uses) and threads them into the forward — -analogous to how Qwen3-Omni threads its 3D-MRoPE cos/sin through preprocess. The -sampling/EOS/logits contract is identical to Orpheus: the non-batched ``forward`` -returns last-token ``logits`` (the KV-cache engine samples them into -``new_token``); ``forward_batched`` samples inside the forward and returns -``new_token`` per request. -""" +"""AR submodule for the Kimi-K2.7 text backbone.""" from __future__ import annotations from typing import Any @@ -43,39 +26,18 @@ class KimiLLMSubmodule(ARNodeSubmodule): - """Autoregressive Kimi/DeepSeek-V3 text backbone (prefill + decode). - - Dispatches on ``graph_walk``: - - ``prefill``: embed the prompt, fill the KV cache, sample the first token; - - ``decode``: embed the previous token, generate the next token. - """ - def __init__(self, language_model: nn.Module, config: KimiK2Config): super().__init__() self.language_model = language_model # KimiForCausalLM self.lm_head = language_model.lm_head self.config = config - # -- CUDA-graph prefill capture grid (full-size defaults) -------------- - # Overridable per-config via ``config.prefill_token_buckets`` / - # ``config.prefill_capture_batch_sizes``: ``KimiK2Config.reduced()`` sets a - # tiny grid for a reduced-size serve, while the full model leaves them - # ``None`` and uses these defaults. Capturing the full 6x5 compiled grid is - # slow, and buckets above a small model's ``max_position_embeddings`` do not - # fit — hence the reduced config trims it. PREFILL_TOKEN_BUCKETS = [32, 64, 128, 256, 512, 1024] PREFILL_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16] def _build_prefill_packed( self, num_tokens: int, device: torch.device, ) -> dict[str, torch.Tensor]: - """Tensor-only post-``preprocess`` packed dict for prefill capture. - - Mirrors ``preprocess``'s tensor outputs: the packed ``(num_tokens,)`` long - ``input_ids`` and the ``(num_tokens,)`` long ``position_ids`` the MLA YARN - RoPE reads. Both are interned as static buffers; at replay the runner - copies the real ``preprocess`` output into them. - """ return { "input_ids": torch.zeros((num_tokens,), dtype=torch.long, device=device), "position_ids": torch.arange(num_tokens, dtype=torch.long, device=device), @@ -84,16 +46,6 @@ def _build_prefill_packed( def get_cuda_graph_configs( self, device: torch.device, tp_world_size: int = 1, ) -> list[BasicBatchedCudaGraphConfig | FlashInferPackedCudaGraphConfig]: - """Decode (per-bs) + prefill (per-token-bucket) captures. - - The YARN-rope path the MLA runs is pure tensor compute (``outer`` + cos/sin - + interleaved rotate) reading ``position_ids`` from a static buffer, so it - captures like Qwen3-Omni's cos/sin-threaded prefill: decode re-runs - ``preprocess`` at replay and copies the packed ``input_ids`` / - ``position_ids`` into the interned buffers; prefill uses the packed dict - above. ``inv_freq`` is lazily cached on first (warmup) forward, so its - address is stable across replay. - """ prefill_buckets = self.config.prefill_token_buckets or self.PREFILL_TOKEN_BUCKETS prefill_batch_sizes = ( self.config.prefill_capture_batch_sizes or self.PREFILL_CAPTURE_BATCH_SIZES @@ -124,8 +76,6 @@ def get_cuda_graph_configs( ), ] - # -- lifecycle --------------------------------------------------------- - def prepare_inputs( self, graph_walk: str, @@ -134,8 +84,6 @@ def prepare_inputs( pos_info: dict[str, PositionInfo] = {}, **kwargs, ) -> ARNodeInputs: - # Cheap host-side: the prompt ids (prefill) or the previous token (decode) - # arrive under the "text_inputs" edge (see KimiK2Model.get_graph_walk_graphs). text_inputs = inputs["text_inputs"][0] return ARNodeInputs( input_ids=text_inputs, @@ -151,17 +99,10 @@ def preprocess( cache_manager = engine_inputs.cache_manager seq_lens = [inp.input_seq_len for inp in inputs] - # Plan attention + rope for the main cache label (CUDA-graph incompatible, - # so it happens here in preprocess, not in forward). cache_manager.set_active_label(_MAIN) cache_manager.plan_attention(seq_lens=seq_lens, is_causal=True, label=_MAIN) cache_manager.plan_rope(seq_lens=seq_lens, pos_ids=None, label=_MAIN) - # Build the per-token YARN position_ids for the MLA rope. These are exactly - # the positions plan_rope uses: each request's position_id_start (advanced - # once per forward by KimiLanguageModel's advance_seq_lens) plus its span. - # request_ids order matches the inputs order (both are batch order), so the - # concatenated input_ids and position_ids stay token-aligned. device = self.get_device() pos_ids_list: list[int] = [] for rid, sl in zip(cache_manager.request_ids, seq_lens, strict=True): @@ -180,12 +121,6 @@ def _hidden( position_ids: torch.Tensor, cache_handle: BatchedCacheManager, ) -> torch.Tensor: - """Token ids -> final hidden states (embed -> layers -> norm). - - ``KimiLanguageModel.forward`` embeds ``input_ids``, runs the decoder stack - (per-layer ``set_layer_idx``, MLA YARN rope from ``position_ids``), calls - ``advance_seq_lens`` once, and returns the normed hidden states. - """ return self.language_model.model(input_ids, cache_handle, position_ids) def forward( @@ -196,8 +131,6 @@ def forward( position_ids: torch.Tensor, **kwargs, ) -> NameToTensorList: - """Non-batched forward: return the last token's logits for the engine to - sample (prefill: last prompt token; decode: the single token).""" cache_handle = engine_inputs.cache_manager hidden = self._hidden(input_ids, position_ids, cache_handle) logits = self.lm_head(hidden[-1:]) @@ -214,10 +147,6 @@ def forward_batched( position_ids: torch.Tensor, **kwargs, ) -> dict[str, NameToTensorList]: - """Batched forward: sample inside the pass (CUDA-graphable sampler) and - return per-request ``new_token``. Prefill packs all requests' tokens, so the - last-token-per-request indices come from the FlashInfer prefill wrapper's - persistent ``qo_indptr`` buffer; decode is one token per request.""" cache_handle = engine_inputs.cache_manager sampler = engine_inputs.sampler cache_handle.set_active_label(_MAIN) @@ -255,8 +184,6 @@ def postprocess( outputs: dict[str, list[torch.Tensor]], **kwargs, ): - # Metadata-only: rebind the new token as the next step's text_inputs. - # EOS is checked in check_stop so the GPU thread doesn't sync on .item(). if "new_token" not in outputs: return outputs["text_inputs"] = outputs["new_token"] diff --git a/mstar/model/kimi_k2_7/weight_loader.py b/mstar/model/kimi_k2_7/weight_loader.py index a7f38e3ab..eb0ce148f 100644 --- a/mstar/model/kimi_k2_7/weight_loader.py +++ b/mstar/model/kimi_k2_7/weight_loader.py @@ -1,33 +1,4 @@ -"""Kimi-K2.7 / DeepSeek-V3 weight loading. - -Maps an HF ``DeepseekV3ForCausalLM`` checkpoint onto the Kimi module tree via the -shared ``load_hf_weights`` machinery (name remap + stacked-shard rules), mirroring -``qwen3_omni_model.py``'s thinker remap/stacked params. - -- :func:`kimi_name_remapper`: strip a ``language_model.`` prefix if present (the - multimodal K2.7-Code repo carries it on its text weights); - ``shared_experts`` -> ``shared_expert``; tag per-routed-expert projections with - an ``__expert{i}__`` marker so one ``shard_id`` carries projection + expert slot. - Vision (``vision_tower.*`` / ``mm_projector.*``) and ``weight_shape`` keys fall - through and are dropped by the base loader's unknown-key skip. -- :func:`build_kimi_stacked_params`: fuse per-expert gate/up -> ``gate_up_proj`` - (w13) and down -> ``down_proj`` (w2); dense/shared gate+up -> ``gate_up_proj``. - Dense rules MUST come after the expert rules — ``_apply_stacked`` returns on - first match and a remapped expert key also contains ``.gate_proj``. -- Router bias (``e_score_correction_bias``) is forced fp32 before load so the - whole-model ``.to(bf16)`` cast can't downcast this fp32 selection bias. - -MLA loads strictly by name — the naive path keeps separate ``q_a_proj`` / -``kv_a_proj_with_mqa``, so no ``fused_qkv_a_proj`` fusion is needed. - -Compressed-tensors INT4 checkpoints: with a ``quant_config`` the stream is -dequantized to bf16 on load (see ``quantization.py``); routed experts can instead -stay packed (``packed_experts=True``) and dequantize inside the fused-expert -kernel. Both are additive — the remap + stacked rules are unchanged. - -Ref: HF key -> param authority is vLLM -``model_executor/models/deepseek_v2.py::DeepseekV2ForCausalLM.load_weights``. -""" +"""HF DeepSeek-V3 checkpoint loading for the Kimi-K2.7 module tree.""" from __future__ import annotations import re @@ -43,50 +14,25 @@ if TYPE_CHECKING: from mstar.model.kimi_k2_7.quantization import CompressedTensorsQuantConfig -# HF suffixes for the per-routed-expert projections. The trailing alternation -# covers a native-bf16 ``.weight`` AND the compressed-tensors sub-keys -# (``.weight_packed`` / ``.weight_scale`` / ``.weight_zero_point``) so a -# packed-expert stream (which passes those sub-keys through raw) is remapped with -# its expert index preserved. A dequant-on-load stream only ever carries ``.weight``. +# Keep the expert index attached while remapping both bf16 and packed sub-keys. _EXPERT_RE = re.compile( r"(.*)\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)" r"\.(weight|weight_packed|weight_scale|weight_zero_point)$" ) -# Base-name matcher (no suffix) for the routed-expert weights kept packed. Used to -# build the ``keep_packed`` predicate handed to the dequant stream. _EXPERT_BASE_RE = re.compile(r"\.experts\.\d+\.(gate_proj|up_proj|down_proj)$") def _is_routed_expert_base(base: str) -> bool: - """True for a routed-expert weight base (``...experts..``). - - ``shared_experts`` does not match — there is no ``.experts..`` (the HF - key is ``mlp.shared_experts.gate_proj``, an underscore not a dotted index), so - the shared expert still dequantizes on load while the routed experts stay packed. - """ return _EXPERT_BASE_RE.search(base) is not None def kimi_name_remapper(name: str) -> str | None: - """HF DeepSeek-V3 checkpoint key -> Kimi module param path. - - Returns ``None`` to drop a key (precomputed ``rotary_emb`` buffers). See the - module docstring for the full mapping; MLA / norms / embed / lm_head are all - identity. Vision (``vision_tower.*`` / ``mm_projector.*``) and ``.weight_shape`` - sub-keys are left unmapped and fall through the base loader's unknown-key skip. - """ if "rotary_emb" in name: return None - # Multimodal K2.7-Code text keys carry a ``language_model.`` prefix; strip it - # only-if-present (a bare ``model.*`` key is left unchanged). if name.startswith("language_model."): name = name[len("language_model."):] - # HF names the shared expert plural; our module has one ``shared_expert``. name = name.replace(".shared_experts.", ".shared_expert.") - # Per-expert fusion marker so the stacked rules can pick up expert index. The - # suffix (``weight`` for bf16, ``weight_packed``/``weight_scale`` for packed - # experts) is carried through so the packed vs bf16 stacked rules can route it. m = _EXPERT_RE.match(name) if m: prefix, expert_idx, proj, suffix = m.groups() @@ -97,20 +43,6 @@ def kimi_name_remapper(name: str) -> str | None: def build_kimi_stacked_params( n_routed_experts: int, packed_experts: bool = False, ) -> list[StackedParamRule]: - """Fused-shard routing for Kimi-K2.7 (mirrors the Qwen3-MoE thinker rules). - - ``packed_experts=False`` (native / dequantized bf16): per-expert ``gate``/``up`` - -> ``experts.gate_up_proj`` (w13) and ``down`` -> ``experts.down_proj`` (w2). - - ``packed_experts=True``: the per-expert ``.weight_packed`` / ``.weight_scale`` - sub-keys route to the FOUR packed params - (``experts.{gate_up_proj,down_proj}_{packed,scale}``), and the bf16 ``.weight`` - expert rules are OMITTED — their ``...__expert{i}__.weight`` source substring - would spuriously match ``...__expert{i}__.weight_packed`` (first-match wins). - - The dense/shared SwiGLU gate/up merge is appended last in both cases (expert - rules precede it so ``.gate_proj`` inside an expert key can't hijack it). - """ rules: list[StackedParamRule] = [] for i in range(n_routed_experts): if packed_experts: @@ -151,19 +83,14 @@ def build_kimi_stacked_params( source_suffix=f".experts.down_proj.__expert{i}__.weight", shard_id=f"down:{i}", )) - # Dense MLP + shared-expert gate/up fusion — AFTER the expert rules. + # Dense/shared gate-up rules must follow expert rules because matching is first-win. rules.append(StackedParamRule(".gate_up_proj", ".gate_proj", 0)) rules.append(StackedParamRule(".gate_up_proj", ".up_proj", 1)) return rules def restore_router_bias_fp32(module: nn.Module) -> None: - """Force every ``e_score_correction_bias`` back to fp32 in place. - - DeepSeek keeps this selection bias fp32; a whole-model ``.to(bfloat16)`` would - downcast it. Call immediately before loading so the source (fp32) copies into - an fp32 destination. - """ + """Force DeepSeek's router selection bias back to fp32 before loading.""" for sub in module.modules(): bias = getattr(sub, "e_score_correction_bias", None) if isinstance(bias, nn.Parameter) and bias.dtype != torch.float32: @@ -177,22 +104,6 @@ def load_kimi_hf_weights( quant_config: "CompressedTensorsQuantConfig | None" = None, packed_experts: bool = False, ) -> set[str]: - """Load an HF DeepSeek-V3 weight stream into ``module``. - - Thin wrapper: restore the fp32 router bias, optionally wrap the stream with the - dequant-on-load parser (when ``quant_config`` is set — the checkpoint is - compressed-tensors quantized), then dispatch through ``load_hf_weights`` with - the Kimi remap + stacked rules. The dequant wrapper emits bf16 ``*.weight`` - keys, so the remap + stacked rules see the same stream as a native-bf16 - checkpoint. Returns the set of param paths that received a tensor (callers can - diff against ``named_parameters()`` to assert completeness). - - ``packed_experts=True``: the routed experts stay PACKED. A ``keep_packed`` - predicate matching routed-expert bases is handed to the dequant stream so those - sub-keys pass through raw (int32 + scale), and the stacked rules route them to - the packed params; every other quantized weight (MLA, dense FFN, shared expert) - still dequantizes to bf16. Requires ``quant_config``. - """ from mstar.model.loader import load_hf_weights if quant_config is not None: @@ -223,12 +134,6 @@ def load_weights( source: str | Path, device: torch.device | str = "cpu", ) -> set[str]: - """``(module, source, device)`` entrypoint mirroring Orpheus. - - ``source`` is a safetensors file or an HF-style checkpoint directory. Picks - the right streaming iterator and drives ``module.load_weights`` (which calls - :func:`load_kimi_hf_weights`). - """ from mstar.model.loader import load_weights as _driver return _driver(module, source, device=device) diff --git a/mstar/utils/flashinfer_utils.py b/mstar/utils/flashinfer_utils.py index a736ca6aa..16028fdf7 100644 --- a/mstar/utils/flashinfer_utils.py +++ b/mstar/utils/flashinfer_utils.py @@ -490,31 +490,11 @@ def set_kv_cache( class FlashInferMLAWrapper: - """Compressed-latent (weight-absorbed) MLA attention over a paged latent cache. - - Wraps ``flashinfer.mla.BatchMLAPagedAttentionWrapper`` (the DeepSeek/Kimi MLA - kernel) for the ``MlaAbsorbCacheManager`` fast path. Unlike the standard - Prefill/Decode wrappers this consumes a **4D latent cache** - ``[max_pages, page_size, head_dim_ckv + head_dim_kpe]`` (one shared MQA latent - per token) and takes the query pre-split into its no-rope / rope parts: - - run(q_nope[T, H, ckv], q_pe[T, H, kpe], - ckv_cache[pages, page_size, ckv], kpe_cache[pages, page_size, kpe]) - -> [T, H, ckv] - - ``ckv``/``kpe`` are passed as **strided views** of the combined latent cache - (``cache[..., :ckv]`` / ``cache[..., ckv:]``) — the kernel accepts them, so the - cache stays a single tensor. - - The kernel is hard-locked to the real Kimi dims (ckv=512, kpe=64): other dims - trigger an *uncatchable* illegal memory access, so the caller must gate on dims - before constructing this (see ``cache_manager._mla_kernel_available``). - - Mirrors the Prefill/Decode wrappers: ``plan`` computes the per-token latent - scatter indices (into static buffers under CUDA graph); ``set_latent`` scatters - ``cat([kv_c, k_pe])`` into the cache; ``run`` calls the kernel. CUDA graph mode - requires ``batch_size``/``max_num_pages`` and static index buffers so ``plan`` - updates values via ``.copy_()`` without reallocating. + """FlashInfer MLA wrapper for the 4D latent cache. + + The kernel only supports real Kimi dims (ckv=512, kpe=64); callers must gate + before construction because off-dim calls can corrupt the CUDA context. CUDA + graph mode updates static index/scatter buffers with ``copy_()``. """ def __init__( @@ -554,8 +534,7 @@ def __init__( assert max_num_pages is not None, "max_num_pages required for CUDA graph mode" assert max_total_tokens is not None, "max_total_tokens required for CUDA graph mode" - # Static index buffers the kernel plan copies into (see the wrapper's - # __init__ docstring): stable addresses across graph replay. + # Stable addresses for graph replay. self._qo_indptr_buf = torch.zeros( batch_size + 1, dtype=torch.int32, device=device ) @@ -604,17 +583,7 @@ def plan( causal: bool = True, dtype: torch.dtype = torch.bfloat16, ): - """Plan the MLA kernel and compute the per-token latent scatter indices. - - Args (all int32, on ``self.device``): - qo_indptr: [n_req + 1] cumulative NEW query tokens per request. - kv_indptr: [n_req + 1] cumulative pages per request. - kv_indices: [total_pages] flattened page indices (per-request order). - kv_len_arr: [n_req] total cached length per request AFTER this append. - - In CUDA graph mode, updates the static index + scatter buffers via - ``.copy_()`` so the captured replay reads stable addresses. - """ + """Plan the MLA kernel and compute latent scatter indices.""" self.dtype = dtype self.attn_wrapper.plan( qo_indptr, @@ -631,10 +600,7 @@ def plan( dtype, ) - # Per-token (page, offset) for the latent scatter, mirroring - # FlashInferPrefillWrapper.plan: token j of request r lands at absolute - # position g = old_len_r + j -> page kv_indices[kv_indptr[r] + g//ps], - # offset g % ps. old_len_r = kv_len_arr[r] - new_tokens_r. + # Map each new token to its latent-cache page/offset from the page table. n_req = qo_indptr.shape[0] - 1 starts = qo_indptr[:-1].to(torch.int32) lens = (qo_indptr[1:] - qo_indptr[:-1]).to(torch.int32) @@ -670,12 +636,7 @@ def plan( @torch.compiler.disable def set_latent(self, latent_cache_layer: torch.Tensor, latent: torch.Tensor): - """Scatter one compressed latent per new token into the paged cache. - - Args: - latent_cache_layer: [max_pages, page_size, ckv + kpe] - latent: [total_tokens, ckv + kpe] = ``cat([kv_c, k_pe])`` for the new tokens. - """ + """Scatter one compressed latent per new token into the paged cache.""" n = self._total_tokens page_idx = self.token_to_page[:n] cache_idx = self.token_to_cache[:n] @@ -689,16 +650,7 @@ def run( ckv_cache: torch.Tensor, kpe_cache: torch.Tensor, ) -> torch.Tensor: - """Run the planned MLA kernel. - - Args: - q_nope: [T, H, ckv] - q_pe: [T, H, kpe] - ckv_cache: [max_pages, page_size, ckv] (strided view of the latent cache) - kpe_cache: [max_pages, page_size, kpe] (strided view of the latent cache) - Returns: - [T, H, ckv] - """ + """Run the planned MLA kernel.""" return self.attn_wrapper.run( q_nope.to(self.dtype), q_pe.to(self.dtype), ckv_cache, kpe_cache, return_lse=False, diff --git a/mstar/utils/fused_moe/kernels.py b/mstar/utils/fused_moe/kernels.py index a8fb61993..4a2d28a8f 100644 --- a/mstar/utils/fused_moe/kernels.py +++ b/mstar/utils/fused_moe/kernels.py @@ -140,7 +140,6 @@ def fused_moe_kernel( @triton.jit def fused_moe_kernel_w4a16( - # Pointers a_ptr, b_ptr, c_ptr, @@ -155,7 +154,6 @@ def fused_moe_kernel_w4a16( K, EM, num_valid_tokens, - # Strides stride_am, stride_ak, stride_be, @@ -169,7 +167,6 @@ def fused_moe_kernel_w4a16( stride_bze, stride_bzk, stride_bzn, - # Block sizes (compile-time) BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, @@ -182,19 +179,10 @@ def fused_moe_kernel_w4a16( HAS_ZP: tl.constexpr, even_Ks: tl.constexpr, ): - """Compute one ``[BLOCK_SIZE_M, BLOCK_SIZE_N]`` output tile from PACKED weights. - - Identical control flow to :func:`fused_moe_kernel`; the only change is the - B load in the K-loop. ``b_ptr`` addresses an int32 tensor of shape - ``(E, N, K // PACK_FACTOR)`` where ``PACK_FACTOR`` INT4 nibbles are packed - low-order-first along the (logical) K axis into each int32. For each K tile we - read the containing int32s, shift out the right nibble, offset-binary subtract - (symmetric: ``- 8``; asymmetric: ``- b_zp``), and scale by the per-``group_size`` - ``b_scale``, all in fp32, then cast to ``compute_type`` for the ``tl.dot``. - - Requires ``BLOCK_SIZE_K % PACK_FACTOR == 0`` and ``BLOCK_SIZE_K % group_size - == 0`` (enforced by :func:`get_default_config`) so a K tile spans a whole - number of int32s and never straddles a group boundary mid-int32. + """Fused MoE tile for packed W4A16 weights. + + ``K`` is logical; weights are int32-packed low-order-first along K. K tiles + must span whole int32s and whole quant groups. """ pid = tl.program_id(axis=0) num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) @@ -219,8 +207,7 @@ def fused_moe_kernel_w4a16( offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N offs_k = tl.arange(0, BLOCK_SIZE_K) a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak) - # Packed B: the int32 holding logical-K index ``kk`` is at ``kk // PACK_FACTOR`` - # along the last (packed) axis; the nibble sits at ``(kk % PACK_FACTOR) * 4``. + # Packed B is indexed by logical K, then shifted to the right INT4 nibble. b_ptrs = ( b_ptr + off_experts * stride_be + (offs_k[:, None] // PACK_FACTOR) * stride_bk @@ -231,8 +218,7 @@ def fused_moe_kernel_w4a16( accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for k_start in range(0, K, BLOCK_SIZE_K): - # Per-group scale index along the logical K axis (same for every int32 in - # a group). Recomputed each tile from the global-K position. + # Group scales are indexed by logical K, not packed K. offs_ks = (offs_k[:, None] + k_start) // group_size b_scale_ptrs = b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn + offs_ks * stride_bsk if even_Ks: @@ -248,13 +234,10 @@ def fused_moe_kernel_w4a16( ) b_packed = tl.load(b_ptrs, mask=k_mask, other=0) b_scale = tl.load(b_scale_ptrs, mask=k_mask, other=1.0).to(tl.float32) - # Extract the nibble. ``>>`` on int32 is arithmetic, but ``& 0xF`` masks - # the sign-extended high bits, so the top nibble (container bit 31 set) is - # exact for all PACK_FACTOR positions. + # Mask after arithmetic shift so the top nibble is exact when bit 31 is set. b_nib = ((b_packed >> b_shifter) & 0xF).to(tl.float32) if HAS_ZP: - # Asymmetric extension (never exercised by Kimi — symmetric only). The - # zero point is stored one-per-group, unpacked, mirroring ``b_scale``. + # Optional asymmetric zero point; Kimi uses symmetric INT4. b_zp_ptrs = b_zp_ptr + off_experts * stride_bze + offs_bn[None, :] * stride_bzn + offs_ks * stride_bzk if even_Ks: b_zp = tl.load(b_zp_ptrs).to(tl.float32) @@ -378,20 +361,10 @@ def invoke_fused_moe_kernel_w4a16( pack_factor: int, group_size: int, ) -> None: - """Launch :func:`fused_moe_kernel_w4a16` (packed-INT4 grouped GEMM). - - Mirrors :func:`invoke_fused_moe_kernel` with two differences that the packed - layout forces: - - * ``K`` is the LOGICAL contraction dim and is passed explicitly — it CANNOT be - derived from ``B_packed.shape[2]`` (that is ``K // pack_factor``). - * ``B_scale`` (shape ``(E, N, K // group_size)``) rides alongside; its strides - are passed ``stride(0), stride(2), stride(1)`` to match the kernel's - ``(bse, bsk, bsn)`` order, exactly like the weight strides. + """Launch the packed-W4A16 MoE kernel. - ``B_zp`` is optional (symmetric INT4 has no zero point). When ``None`` the - kernel's ``HAS_ZP`` is off and a stand-in tensor (``B_scale``) is passed so the - dead zp-stride args stay valid; nothing is dereferenced. + ``K`` is logical, while ``B_packed.shape[2]`` is ``K // pack_factor``. When + ``B_zp`` is absent, ``HAS_ZP`` gates all loads from the stand-in tensor. """ assert topk_weights.stride(1) == 1 assert sorted_token_ids.stride(0) == 1 @@ -635,12 +608,9 @@ def get_default_config( For decode batch sizes (``M`` on the order of 1--64) we always fall into the ``M <= E`` branch since Qwen3-Omni has ``E == 128``. - ``group_size`` (set only on the W4A16 path) does NOT change the bf16 result: - when ``None`` the returned config is byte-for-byte the historical one. When - set, ``BLOCK_SIZE_K`` is clamped down (halving) until it is a multiple of both - ``pack_factor`` (8 for INT4) and ``group_size`` — the divisibility the packed - kernel needs so a K tile spans whole int32s and whole groups. Kimi - (``group_size=32``; configs 64/32) already complies, so the clamp is a no-op. + When ``group_size`` is set for W4A16, ``BLOCK_SIZE_K`` is clamped until it is + divisible by both INT4 pack factor and group size. ``None`` preserves the + historical config. """ if M <= E: config = { diff --git a/mstar/utils/fused_moe/runner.py b/mstar/utils/fused_moe/runner.py index 66d98bd16..b519d5ea0 100644 --- a/mstar/utils/fused_moe/runner.py +++ b/mstar/utils/fused_moe/runner.py @@ -55,13 +55,11 @@ def fused_experts( ``(num_experts, 2 * moe_intermediate_size, hidden)``. Matches the ``experts.gate_up_proj`` parameter the WeightConverter already produces in :mod:`mstar.model.qwen3_omni.qwen3_omni_model`. - On the W4A16 path (``w1_scale`` set) this is instead the PACKED int32 - tensor ``(num_experts, 2 * moe_intermediate_size, hidden // pack_factor)``. + Packed int32 on the W4A16 path. w2 : torch.Tensor Down projection weights, shape ``(num_experts, hidden, moe_intermediate_size)``. Matches - ``experts.down_proj``. On the W4A16 path this is the packed int32 tensor - ``(num_experts, hidden, moe_intermediate_size // pack_factor)``. + ``experts.down_proj``. Packed int32 on the W4A16 path. topk_weights : torch.Tensor ``(tokens, top_k)``, routing probabilities (possibly renormalized). Dtype matches ``hidden_states``. @@ -75,14 +73,11 @@ def fused_experts( ``(tokens, top_k, hidden)`` — the caller is responsible for the reduce (e.g. after an all-reduce for TP). w1_scale, w2_scale : torch.Tensor | None - W4A16 group scales, shapes ``(E, 2*inter, hidden//group_size)`` - and ``(E, hidden, inter//group_size)``. When BOTH are ``None`` (default) - the bf16 path runs byte-for-byte as before; when set, ``w1``/``w2`` are - packed int32 and the in-kernel dequant path runs. + W4A16 group scales; both ``None`` keeps the historical bf16 path. w1_zp, w2_zp : torch.Tensor | None Optional asymmetric zero points (unused for Kimi's symmetric INT4). group_size, pack_factor : int | None - Group granularity and INT4 pack factor (8); required on the W4A16 path. + Required on the W4A16 path. Returns ------- @@ -93,7 +88,7 @@ def fused_experts( assert hidden_states.is_contiguous(), "hidden_states must be contiguous" assert hidden_states.dim() == 2 assert topk_weights.shape == topk_ids.shape - # Activations stay bf16/fp16 on both paths — only the WEIGHTS are quantized. + # Only weights are quantized; activations stay bf16/fp16. assert hidden_states.dtype in (torch.bfloat16, torch.float16) quantized = w1_scale is not None @@ -159,7 +154,6 @@ def fused_experts( ) # 3. Gate+up GEMM: cache1[slot] = hidden[slot // top_k] @ w1[expert].T - # (W4A16: w1 packed int32, dequantized in-kernel; GEMM-1 contracts over hidden.) if quantized: invoke_fused_moe_kernel_w4a16( A=hidden_states, @@ -202,7 +196,6 @@ def fused_experts( # 5. Down GEMM (weighted): cache3[slot] = topk_weight[slot] * (cache2[slot] @ w2[expert].T) # top_k=1 for this GEMM so the kernel's offs_token // top_k is identity # -- it reads cache2 rows directly instead of the (slot // top_k)-th source row. - # (W4A16: w2 packed int32; GEMM-2 contracts over the intermediate dim.) if quantized: invoke_fused_moe_kernel_w4a16( A=cache2, diff --git a/mstar/utils/marlin/__init__.py b/mstar/utils/marlin/__init__.py index fdcd72723..74a404c88 100644 --- a/mstar/utils/marlin/__init__.py +++ b/mstar/utils/marlin/__init__.py @@ -1,10 +1,4 @@ -"""Vendored Marlin W4A16 (INT4) CUDA kernels for mstar. - -JIT-compiled from Apache-2.0 vLLM sources under ``csrc/`` on first use, with a -Triton fallback (see :mod:`mstar.utils.marlin.loader`). Exposes the repack + GEMM -launchers (:mod:`mstar.utils.marlin.ops`) used by the compressed-tensors W4A16 -routed-expert path. -""" +"""Vendored Marlin W4A16 CUDA kernels for mstar.""" from mstar.utils.marlin.loader import is_marlin_available __all__ = ["is_marlin_available"] diff --git a/mstar/utils/marlin/loader.py b/mstar/utils/marlin/loader.py index 2d126b0bd..db72a5321 100644 --- a/mstar/utils/marlin/loader.py +++ b/mstar/utils/marlin/loader.py @@ -1,16 +1,4 @@ -"""JIT build + load of the vendored Marlin W4A16 CUDA ops. - -Mirrors :mod:`mstar.utils.fused_moe.align`: the Marlin CUDA sources (vendored -Apache-2.0 from vLLM under ``csrc/``) are JIT-compiled with -``torch.utils.cpp_extension.load`` on first use and registered as -``torch.ops._mstar_marlin_C.*`` — no ``vllm`` / ``sgl_kernel`` runtime dependency. - -If the build fails (no ``nvcc`` / no ``ninja`` / sm<80 / ABI mismatch) the loader -logs and returns ``False``; callers fall back to the Triton W4A16 path -(``fused_moe_kernel_w4a16``). Marlin is a *speed* layer over a correctness path -that already exists — exactly the CUDA-op-or-torch-fallback pattern ``align.py`` -uses for ``moe_align_block_size``. -""" +"""JIT build + load of the vendored Marlin W4A16 CUDA ops.""" from __future__ import annotations import functools @@ -25,10 +13,6 @@ _MARLIN = os.path.join(_CSRC, "libtorch_stable", "quantization", "marlin") _MARLIN_MOE = os.path.join(_CSRC, "libtorch_stable", "moe", "marlin_moe_wna16") -# Sources compiled into the ``_mstar_marlin_C`` extension: the repack op, the MoE -# GEMM host+device shim, and the pre-generated per-config kernel instantiations -# (``sm80_kernel_*.cu`` + ``kernel_selector.h``, produced once by the trimmed -# ``generate_kernels.py`` and vendored — GPTQ symmetric INT4, fp16/bf16 only). _SOURCES = [ os.path.join(_MARLIN, "gptq_marlin_repack.cu"), os.path.join(_MARLIN_MOE, "marlin_moe.cu"), @@ -36,18 +20,11 @@ os.path.join(_MARLIN_MOE, "sm80_kernel_float16_u4b8_float16.cu"), ] -# Marlin's device code (cp.async, m16n8k16 MMA, bf16) requires sm80+. _MIN_CAPABILITY = (8, 0) @functools.lru_cache(maxsize=1) def is_marlin_available() -> bool: - """JIT-build the Marlin ops once per process; return whether they are usable. - - Cached so compilation is attempted at most once. Any failure (missing - toolchain, unsupported GPU, compile/ABI error) is logged and the caller uses - the Triton W4A16 fallback. - """ if not torch.cuda.is_available(): return False capability = torch.cuda.get_device_capability() @@ -69,7 +46,6 @@ def is_marlin_available() -> bool: extra_cuda_cflags=["-O3", "-std=c++17", "--expt-relaxed-constexpr"], verbose=False, ) - # Touch an op so a registration failure surfaces here, not at call time. _ = torch.ops._mstar_marlin_C.moe_wna16_marlin_gemm return True except Exception as e: # pragma: no cover -- depends on the build toolchain diff --git a/mstar/utils/marlin/ops.py b/mstar/utils/marlin/ops.py index c01988384..814e0a72c 100644 --- a/mstar/utils/marlin/ops.py +++ b/mstar/utils/marlin/ops.py @@ -1,16 +1,4 @@ -"""Python launchers over the vendored ``torch.ops._mstar_marlin_C`` Marlin ops. - -Mirrors vLLM's ``_custom_ops`` + ``marlin_utils`` helpers (same transformation -sequence: checkpoint INT4 → Marlin-repacked → GEMM), so the port is auditable -against the reference. Callers must gate on -:func:`mstar.utils.marlin.is_marlin_available` first — these dereference the JIT -op namespace directly. - -The routed-expert GEMM (:func:`fused_marlin_moe`) reuses mstar's existing MoE -plumbing — ``moe_align_block_size`` (token→expert sort), ``act_and_mul_triton`` -(SwiGLU), and ``moe_sum_reduce_triton`` (top-k fold) — so only the two INT4 -matmuls are Marlin; everything around them is shared with the bf16/Triton paths. -""" +"""Python launchers over the vendored Marlin torch ops.""" from __future__ import annotations import torch @@ -19,23 +7,12 @@ from mstar.utils.fused_moe.kernels import act_and_mul_triton, moe_sum_reduce_triton from mstar.utils.marlin.scalar_type import UINT4B8_ID -# Marlin repacked-weight tile (from the vendored marlin.cuh ``tile_size``). _MARLIN_TILE = 16 -# --------------------------------------------------------------------------- -# Load-time repack (checkpoint GPTQ-packed INT4 → Marlin tiled layout) -# --------------------------------------------------------------------------- - def gptq_marlin_repack( b_q_weight: torch.Tensor, size_k: int, size_n: int, num_bits: int = 4 ) -> torch.Tensor: - """Repack a GPTQ-layout packed INT4 weight into Marlin tiled layout. - - ``b_q_weight`` is int32 ``(size_k // pack_factor, size_n)`` (K-major packed, - the layout Marlin's repack expects). Returns the Marlin-tiled int32 weight - ``(size_k // 16, size_n * 16 // pack_factor)``. - """ perm = torch.empty(0, dtype=torch.int32, device=b_q_weight.device) return torch.ops._mstar_marlin_C.gptq_marlin_repack( b_q_weight, perm, size_k, size_n, num_bits @@ -45,11 +22,6 @@ def gptq_marlin_repack( def gptq_marlin_moe_repack( b_q_weight: torch.Tensor, size_k: int, size_n: int, num_bits: int = 4 ) -> torch.Tensor: - """Per-expert :func:`gptq_marlin_repack` (mirrors vLLM's pure-Python loop). - - ``b_q_weight`` is int32 ``(E, size_k // pack_factor, size_n)``; returns - ``(E, size_k // 16, size_n * num_bits // 8)``. - """ num_experts = b_q_weight.shape[0] perm = torch.empty(0, dtype=torch.int32, device=b_q_weight.device) output = torch.empty( @@ -65,7 +37,6 @@ def gptq_marlin_moe_repack( def _get_scale_perms() -> tuple[list[int], list[int]]: - """Marlin scale-permutation index tables (verbatim from vLLM marlin_utils).""" scale_perm: list[int] = [] for i in range(8): scale_perm.extend([i + 8 * j for j in range(8)]) @@ -78,7 +49,6 @@ def _get_scale_perms() -> tuple[list[int], list[int]]: def marlin_permute_scales( s: torch.Tensor, size_k: int, size_n: int, group_size: int ) -> torch.Tensor: - """Permute a single expert's group scales into Marlin layout (vLLM parity).""" scale_perm, scale_perm_single = _get_scale_perms() if group_size < size_k and group_size != -1: s = s.reshape((-1, len(scale_perm)))[:, scale_perm] @@ -90,7 +60,6 @@ def marlin_permute_scales( def marlin_moe_permute_scales( s: torch.Tensor, size_k: int, size_n: int, group_size: int ) -> torch.Tensor: - """Per-expert :func:`marlin_permute_scales`. ``s`` is ``(E, num_groups, size_n)``.""" num_experts = s.shape[0] output = torch.empty_like(s) for e in range(num_experts): @@ -99,15 +68,10 @@ def marlin_moe_permute_scales( def marlin_make_workspace(device: torch.device, max_blocks_per_sm: int = 4) -> torch.Tensor: - """Marlin reduce/lock workspace: one int per (SM × max_blocks_per_sm).""" sms = torch.cuda.get_device_properties(device).multi_processor_count return torch.zeros(sms * max_blocks_per_sm, dtype=torch.int, device=device) -# --------------------------------------------------------------------------- -# Runtime: fused routed-expert Marlin GEMM -# --------------------------------------------------------------------------- - def fused_marlin_moe( hidden_states: torch.Tensor, w1_marlin: torch.Tensor, @@ -121,15 +85,6 @@ def fused_marlin_moe( activation: str = "silu", reduce_results: bool = True, ) -> torch.Tensor: - """Marlin W4A16 routed-expert dispatch (gate_up GEMM → SwiGLU → down GEMM). - - Layout mirrors :func:`mstar.utils.fused_moe.fused_experts` but with Marlin - kernels: ``w1_marlin``/``w2_marlin`` are the Marlin-repacked int32 experts - (from :func:`gptq_marlin_moe_repack`), ``w1_scale``/``w2_scale`` the permuted - group scales (from :func:`marlin_moe_permute_scales`). Returns - ``(tokens, hidden)`` when ``reduce_results`` else the per-slot - ``(tokens, top_k, hidden)`` tensor the TP path all-reduces before folding. - """ assert hidden_states.is_contiguous() and hidden_states.dim() == 2 assert hidden_states.dtype in (torch.bfloat16, torch.float16) M, K = hidden_states.shape diff --git a/mstar/utils/marlin/scalar_type.py b/mstar/utils/marlin/scalar_type.py index b1a9c3e5f..c1edc41cd 100644 --- a/mstar/utils/marlin/scalar_type.py +++ b/mstar/utils/marlin/scalar_type.py @@ -1,18 +1,7 @@ -"""vLLM ``ScalarType`` ids mirrored in Python for the Marlin ops. - -The vendored ``core/scalar_type.hpp`` packs a scalar type into a single int64 id -(field order exponent|mantissa|signed|bias|finite|nan_repr). The Marlin GEMM op -takes ``b_type_id`` and reconstructs the type via ``ScalarType::from_id``, so the -Python side must pass the exact same id. Only symmetric INT4 (``uint4b8``) is -built. The id is computed directly from the header's bit layout and re-validated -at runtime by the C++ op's ``TORCH_CHECK(b_type == kU4B8)``. - - uint4b8 = ScalarType::uint(size_bits=4, bias=8) - = (mantissa=4 << 8) | (bias=8 << 17) | (nan_repr=NAN_IEEE_754=1 << 50) -""" +"""vLLM ``ScalarType`` ids mirrored in Python for the Marlin ops.""" from __future__ import annotations -# vllm::kU4B8.id() — symmetric GPTQ-style INT4 (offset-binary, subtract 8). +# vllm::kU4B8.id(): (mantissa=4 << 8) | (bias=8 << 17) | (nan_repr=1 << 50). UINT4B8_ID = (4 << 8) | (8 << 17) | (1 << 50) # == 1125899907892224 assert UINT4B8_ID == 1125899907892224, "uint4b8 id drifted from the vendored header" diff --git a/pyproject.toml b/pyproject.toml index 0fb9d1b03..1467f82fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -205,8 +205,6 @@ exclude = ["tests*", "benchmark*", "examples*"] # Ship the vendored CUDA source so it can be JIT-compiled on first use. [tool.setuptools.package-data] "mstar.utils.fused_moe" = ["csrc/*.cu", "csrc/*.cuh", "csrc/*.h", "csrc/*.cpp"] -# Marlin W4A16 sources live in a nested csrc tree (core/, libtorch_stable/...), -# including the pre-generated per-config kernel instantiations. "mstar.utils.marlin" = [ "csrc/**/*.cu", "csrc/**/*.cuh", "csrc/**/*.h", "csrc/**/*.hpp", "csrc/**/*.cpp", "csrc/**/*.py", diff --git a/test/integration/test_kimi_components.py b/test/integration/test_kimi_components.py index 8cf9d6a3a..4a29a3d69 100644 --- a/test/integration/test_kimi_components.py +++ b/test/integration/test_kimi_components.py @@ -1,17 +1,3 @@ -"""M1 golden tests for Kimi-K2.7 cheap reused components. - -For each cheap component of the DeepSeek-V3 text backbone — RMSNorm, the dense -SwiGLU MLP, the token embedding, and the LM head — build the mstar component from -``KimiK2Config.reduced()``, load identical random weights into it and an -independent reference, and assert the outputs match. The reference formulas are -inlined here (self-contained; no dependency on the local golden harness) and each -is cited to the vLLM DeepSeek-V3 source so the golden is authoritative. - -This is a GPU test: mstar's standard RMSNorm dispatches to a FlashInfer fused -kernel, so the suite runs on ``cuda``. It skips automatically without a GPU. - -Run: pytest test/integration/test_kimi_components.py -v -""" import pytest import torch import torch.nn.functional as F @@ -36,14 +22,8 @@ def _cfg() -> KimiK2Config: return KimiK2Config.reduced() -# -------------------------------------------------------------------------- -# Independent references (cited to vllm-project/vllm .../models/deepseek_v2.py) -# -------------------------------------------------------------------------- - def _ref_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: - # Standard Llama/DeepSeek RMSNorm (HF ``LlamaRMSNorm`` / vLLM - # ``RMSNorm.forward_native``): normalize in fp32, scale by weight in the - # input dtype. + # DeepSeek/Llama RMSNorm: normalize in fp32, scale in input dtype. orig_dtype = x.dtype x32 = x.float() var = x32.pow(2).mean(dim=-1, keepdim=True) @@ -54,8 +34,7 @@ def _ref_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Ten def _ref_swiglu( x: torch.Tensor, gate_w: torch.Tensor, up_w: torch.Tensor, down_w: torch.Tensor ) -> torch.Tensor: - # ``DeepseekV2MLP.forward`` = ``down_proj(SiluAndMul(gate_up_proj(x)))``, - # bias=False, silu-only. ``SiluAndMul([g, u]) = silu(g) * u``. + # DeepseekV2MLP uses bias-free down_proj(silu(gate) * up). gate = F.linear(x, gate_w) up = F.linear(x, up_w) return F.linear(F.silu(gate) * up, down_w) @@ -66,14 +45,9 @@ def _ref_embedding(ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: def _ref_lm_head(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: - # Untied head (``tie_word_embeddings=False``): plain ``x @ weight.T``. return F.linear(x, weight) -# -------------------------------------------------------------------------- -# Tests -# -------------------------------------------------------------------------- - def test_rmsnorm_matches_reference(): torch.manual_seed(0) cfg = _cfg() @@ -102,7 +76,6 @@ def test_dense_mlp_matches_reference(): gate_w = torch.randn(i, h, device=DEVICE, dtype=dtype) * 0.05 up_w = torch.randn(i, h, device=DEVICE, dtype=dtype) * 0.05 down_w = torch.randn(h, i, device=DEVICE, dtype=dtype) * 0.05 - # Load through the real (stacked) loaders: gate -> shard 0, up -> shard 1. mlp.gate_up_proj.weight_loader(mlp.gate_up_proj.weight, gate_w, loaded_shard_id=0) mlp.gate_up_proj.weight_loader(mlp.gate_up_proj.weight, up_w, loaded_shard_id=1) mlp.down_proj.weight_loader(mlp.down_proj.weight, down_w) diff --git a/test/integration/test_kimi_decoder_layer.py b/test/integration/test_kimi_decoder_layer.py index d4b5598cf..f7b75f750 100644 --- a/test/integration/test_kimi_decoder_layer.py +++ b/test/integration/test_kimi_decoder_layer.py @@ -1,22 +1,3 @@ -"""M4 golden tests for the Kimi-K2.7 / DeepSeek-V3 decoder layer. - -One golden per feed-forward variant — a dense layer (``layer_idx=0``, below -``first_k_dense_replace``) and a MoE layer (``layer_idx=1``) — each compared to a -self-contained inline reference that re-derives the whole block: -pre-norm → naive-MLA self-attention → residual → pre-norm → (dense-or-MoE) FFN → -residual. The inner attention and FFN references are the same ones the M2/M3 -goldens use, cited to vLLM; what this test adds is the residual/norm *wiring*, -matching vLLM ``models/deepseek_v2.py::DeepseekV2DecoderLayer.forward``. - -A ``_MockMLACache`` stands in for the paged cache (causal SDPA at the fixed -``1/sqrt(qk_head_dim)`` scale FlashInfer uses); the real paged ``run_attention`` -is exercised separately in ``test_kimi_flashinfer_attention.py``. - -GPU test (mstar RMSNorm + the fused expert GEMM are CUDA/half-precision only); -skips without a GPU. - -Run: pytest test/integration/test_kimi_decoder_layer.py -v -""" import pytest import torch import torch.nn.functional as F @@ -39,11 +20,6 @@ DEVICE = "cuda" -# -------------------------------------------------------------------------- -# Inline references (cited to vLLM deepseek_v2.py / deepseek_scaling_rope.py / -# cpu_fused_moe.py) — self-contained, no dependency on the golden harness. -# -------------------------------------------------------------------------- - def _ref_rmsnorm(x, weight, eps): x32 = x.float() x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) @@ -67,7 +43,6 @@ def _ref_yarn_rope(pos, q_pe, k_pe, rotary_dim, base, factor, max_pos, def _sdpa_causal(q, k, v, scale): - """Causal SDPA at a fixed scale (mirrors _MockMLACache / FlashInfer).""" qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) # (H,T,D) T = q.shape[0] causal = torch.triu( @@ -77,7 +52,6 @@ def _sdpa_causal(q, k, v, scale): def _ref_attn_forward(attn, cfg, h_normed, pos): - """Independent naive-MLA forward matching KimiMLAAttention.forward.""" T, H = h_normed.shape[0], attn.num_heads Dnope, Drope, Dv, L = ( cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank) @@ -146,7 +120,6 @@ def _ref_swiglu(x, gate_w, up_w, down_w): def _ref_mlp_forward(mlp, cfg, h_normed): - """Dense SwiGLU or MoE (routed + ungated shared), matching the module.""" if isinstance(mlp, KimiSparseMoeBlock): I = cfg.moe_intermediate_size si = cfg.moe_intermediate_size * cfg.n_shared_experts @@ -176,11 +149,6 @@ def _ref_decoder_layer(layer, cfg, h, pos): return h1 + _ref_mlp_forward(layer.mlp, cfg, mlp_in) -# -------------------------------------------------------------------------- -# Mock paged cache: causal SDPA at 1/sqrt(head_dim), no cross-layer history -# (a single prefill forward; each layer attends its own q/k/v). -# -------------------------------------------------------------------------- - class _MockMLACache: def __init__(self, head_dim): self.scale = head_dim ** -0.5 @@ -206,7 +174,7 @@ def _build_layer(cfg, layer_idx, dtype): layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02) mlp = layer.mlp if isinstance(mlp, KimiSparseMoeBlock): - # Keep the router fp32 (deterministic selection); experts/shared bf16. + # Keep the router fp32 for deterministic selection. mlp.gate.weight.data = torch.randn( cfg.n_routed_experts, cfg.hidden_size, device=DEVICE) mlp.gate.e_score_correction_bias.data = torch.randn( @@ -221,10 +189,6 @@ def _build_layer(cfg, layer_idx, dtype): return layer -# -------------------------------------------------------------------------- -# Tests -# -------------------------------------------------------------------------- - def test_dense_decoder_layer_matches_reference(): torch.manual_seed(0) cfg = KimiK2Config.reduced() diff --git a/test/integration/test_kimi_flashinfer_attention.py b/test/integration/test_kimi_flashinfer_attention.py index c4bb9180c..aed3e20c7 100644 --- a/test/integration/test_kimi_flashinfer_attention.py +++ b/test/integration/test_kimi_flashinfer_attention.py @@ -1,35 +1,3 @@ -"""M4 step 3: validate mstar's REAL paged ``run_attention`` for the naive MLA. - -The naive/materialized MLA stores a ``head_dim = qk_head_dim`` (nope+rope) K plus -a V padded to that same width, then calls the paged ``run_attention`` at the -fixed ``1/sqrt(head_dim)`` scale FlashInfer uses (which is exactly why the -``mscale^2`` softmax boost is folded into q — ``run_attention`` exposes no custom -``sm_scale``). This test drives the **real** ``FlashInferCacheManager`` over a -genuine paged KV cache (real ``PagedAllocationManager`` + ``LocalTransferEngine``) -and asserts its ``run_attention`` matches a causal-SDPA reference at -``1/sqrt(head_dim)`` — confirming both the paged path integrates and the scale -assumption the naive MLA relies on. - -KEY CONSTRAINT FOUND (this is the "FlashInfer-192" answer, recorded for M5/M6): -FlashInfer 0.6.14's SM90 (Hopper / H200) prefill kernel has a compile-time -``static_assert(HEAD_DIM_VO == 64 || HEAD_DIM_VO == 128 || HEAD_DIM_VO == 256)`` -(``flashinfer/.../attention/hopper/prefill_sm90.cuh:572``). The naive MLA pads V -to ``qk_head_dim``, so ``head_dim_vo == head_dim``: - - * real Kimi ``qk_head_dim = 192`` (nope 128 + rope 64) -> vo=192 -> JIT FAILS - * reduced-config ``qk_head_dim = 24`` -> vo=24 -> JIT FAILS - * 64 / 128 / 256 -> supported -> OK - -So the naive MLA path cannot use the paged ``run_attention`` at head_dim 192 (or -the reduced 24) on Hopper as-is. Validated mitigation (M6 follow-up, NOT done -here): pad ``head_dim`` up to the next supported vo (256 for real 192, 64 for the -reduced 24), pad q/k/v to it, and slice the attention output back to -``v_head_dim``. The supported-dim runs below (128 / 256) are exactly that padded -path; the env-gated test at the bottom records the raw 192 failure. - -Run: pytest test/integration/test_kimi_flashinfer_attention.py -v - KIMI_TEST_FLASHINFER_192=1 pytest ... -k rejects # ~60s failing JIT -""" import os import pytest @@ -52,14 +20,6 @@ def _make_real_cache_manager(num_heads, head_dim, dtype, page_size=128, max_num_pages=8): - """Build a genuine paged FlashInferCacheManager for one request. - - Mirrors ``KVCacheEngine.load_model`` / ``_create_cache_manager``: a real - ``[layers, pages, 2, page_size, heads, head_dim]`` KV cache, a - ``PagedAllocationManager`` over a no-op ``LocalTransferEngine`` (single-node - SHM path — no cross-worker reads), and the flashinfer backend. Returns - ``(cache_manager, alloc_manager)`` so the caller can clean up. - """ kv_cache = torch.zeros( 2, max_num_pages, 2, page_size, num_heads, head_dim, dtype=dtype, device=DEVICE, @@ -100,13 +60,8 @@ def _sdpa_causal(q, k, v, scale): return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) -# FlashInfer SM90 prefill supports head_dim_vo in {64, 128, 256}. These are the -# sizes the naive-MLA V-pad would target: 64 for the reduced config (24 -> 64), -# 128 canonical, 256 for the real Kimi qk_head_dim (192 -> 256). @pytest.mark.parametrize("head_dim", [128, 256]) def test_real_paged_run_attention_matches_sdpa(head_dim): - """The real FlashInferCacheManager.run_attention == causal SDPA at - 1/sqrt(head_dim), the fixed scale the naive MLA folds mscale^2 into q for.""" torch.manual_seed(0) num_heads, T = 4, 6 dtype = torch.bfloat16 @@ -135,15 +90,6 @@ def test_real_paged_run_attention_matches_sdpa(head_dim): "the head_dim=192 SM90 static_assert rejection", ) def test_flashinfer_rejects_head_dim_192(): - """Executable record of the constraint: the real paged run_attention cannot - JIT-build for head_dim=192 (vo=192) on Hopper — FlashInfer static_asserts - HEAD_DIM_VO in {64,128,256}. If FlashInfer/the mitigation ever lifts this, - this test flips to failing (no exception raised) and flags the change. - - The failure surfaces as the JIT build erroring out; the concrete exception - type varies by stage (a RuntimeError wrapping the ninja/nvcc - CalledProcessError), so we assert on the broad base and check the message - points at the build rather than an unrelated error.""" torch.manual_seed(0) num_heads, T, head_dim = 4, 6, 192 dtype = torch.bfloat16 @@ -153,16 +99,13 @@ def test_flashinfer_rejects_head_dim_192(): k = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1 v = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1 cm.set_active_label("main") - # The offending kernel is JIT-built when FlashInfer schedules it — that - # can happen in plan_attention (the wrapper.plan() call) or run_attention - # depending on version, so both sit inside the raises block. + # FlashInfer may JIT the failing kernel in plan_attention or run_attention. with pytest.raises(Exception) as exc_info: cm.plan_attention(seq_lens=[T], is_causal=True, dtype=dtype) cm.set_layer_idx(0) cm.run_attention(q=q, k=k, v=v) torch.cuda.synchronize() - # Guard against catching an unrelated error: the message must reference - # the failed build / the offending head_dim. + # Guard against catching an unrelated error. msg = str(exc_info.value).lower() assert "ninja" in msg or "build" in msg or "192" in msg finally: diff --git a/test/integration/test_kimi_forward.py b/test/integration/test_kimi_forward.py index ce426b940..98de088aa 100644 --- a/test/integration/test_kimi_forward.py +++ b/test/integration/test_kimi_forward.py @@ -1,27 +1,3 @@ -"""M4 full-forward golden test for Kimi-K2.7 / DeepSeek-V3 (assembled backbone). - -Runs ``KimiForCausalLM`` end to end on the reduced config — token ids → -embedding → stacked ``KimiDecoderLayer`` blocks (including the dense→MoE -transition at ``first_k_dense_replace=1``) → final RMSNorm → untied LM head → -logits — and compares against a self-contained inline reference that re-derives -every step. The inner attention / FFN / router references are the same ones the -M2/M3 goldens use (cited to vLLM); this test verifies the *assembly*: embedding, -the per-layer cache-handle contract (``set_layer_idx`` each layer, -``advance_seq_lens`` once after), the layer stack, and the LM head. - -A ``_MockMLACache`` stands in for the paged cache (causal SDPA at the fixed -``1/sqrt(qk_head_dim)`` scale). Because a single prefill forward re-attends the -same tokens each layer with no cross-layer history, the mock — which attends the -q/k/v of each ``run_attention`` call independently — reproduces the paged -prefill exactly. The real FlashInfer paged path is validated separately in -``test_kimi_flashinfer_attention.py``. - -Refs: vLLM ``models/deepseek_v2.py`` (``DeepseekV2Model`` / ``DecoderLayer`` / -``DeepseekV2MoE``), ``rotary_embedding/deepseek_scaling_rope.py``, -``fused_moe/cpu_fused_moe.py::grouped_topk``. - -Run: pytest test/integration/test_kimi_forward.py -v -""" import pytest import torch import torch.nn.functional as F @@ -44,11 +20,6 @@ DEVICE = "cuda" -# -------------------------------------------------------------------------- -# Inline references (self-contained; cited to vLLM). Same math as the M2/M3 -# component goldens, assembled here into a whole-model forward. -# -------------------------------------------------------------------------- - def _ref_rmsnorm(x, weight, eps): x32 = x.float() x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) @@ -99,9 +70,7 @@ def _ref_attn_forward(attn, cfg, h_normed, pos): pos, q_pe, k_pe, Drope, cfg.rope_theta, r["factor"], r["original_max_position_embeddings"], r.get("beta_fast", 32), r.get("beta_slow", 1), r.get("mscale", 1.0), r.get("mscale_all_dim", 0.0)) - # M6 mitigation: q/k padded from Dqk and v from Dv up to padded_head_dim; the - # softmax_scale_boost compensates so run_attention's padded_head_dim**-0.5 - # scale reproduces the DeepSeek qk_head_dim**-0.5 * mscale**2 scale. + # softmax_scale_boost restores DeepSeek's scale after padded-head attention. pad = cfg.padded_head_dim q = F.pad(torch.cat([q_nope, q_pe], dim=-1), [0, pad - cfg.qk_head_dim]) * attn.softmax_scale_boost k = F.pad(torch.cat([k_nope, k_pe.expand(T, H, Drope)], dim=-1), [0, pad - cfg.qk_head_dim]) @@ -189,10 +158,6 @@ def _ref_forward(model, cfg, ids, pos): return F.linear(h, model.lm_head.weight) -# -------------------------------------------------------------------------- -# Mock paged cache (causal SDPA at 1/sqrt(head_dim)) + weight init -# -------------------------------------------------------------------------- - class _MockMLACache: def __init__(self, head_dim): self.scale = head_dim ** -0.5 @@ -242,16 +207,11 @@ def _build_model(cfg, dtype): return model.eval() -# -------------------------------------------------------------------------- -# Test -# -------------------------------------------------------------------------- - def test_full_forward_logits_match_reference(): torch.manual_seed(0) cfg = KimiK2Config.reduced() dtype = torch.bfloat16 model = _build_model(cfg, dtype) - # The stack spans the dense->MoE transition (first_k_dense_replace=1). assert not isinstance(model.model.layers[0].mlp, KimiSparseMoeBlock) assert isinstance(model.model.layers[1].mlp, KimiSparseMoeBlock) @@ -264,8 +224,7 @@ def test_full_forward_logits_match_reference(): got = model(ids, cache, pos) expected = _ref_forward(model, cfg, ids, pos) - # advance_seq_lens is called exactly once per forward (after the layer loop), - # not once per layer — the cache-handle contract mirrored from Orpheus. + # advance_seq_lens belongs after the layer loop, not once per layer. assert cache.advance_calls == 1 assert got.shape == (T, cfg.vocab_size) torch.testing.assert_close(got, expected, rtol=3e-2, atol=3e-2) diff --git a/test/integration/test_kimi_mla.py b/test/integration/test_kimi_mla.py index 602e9d8b7..347796584 100644 --- a/test/integration/test_kimi_mla.py +++ b/test/integration/test_kimi_mla.py @@ -1,22 +1,3 @@ -"""M3 golden tests for Kimi-K2.7 MLA attention (naive/materialized path). - -Three goldens against independent references cited to vLLM: - - YARN RoPE (KimiYarnRotaryEmbedding) vs a DeepseekScalingRotaryEmbedding-style - forward_static, - - the q/k/v assembly (projections + rope-on-slice + k_pe broadcast + v-pad + - mscale^2 q-prescale) captured at ``run_attention`` via a mock cache handle, and - - the full attention forward (+ causal attention + output slice + o_proj). - -A ``_MockMLACache`` stands in for the paged cache: its ``run_attention`` does a -causal SDPA at the fixed ``1/sqrt(qk_head_dim)`` scale (what FlashInfer uses), -which is what lets us golden the MLA math without the paged engine. The real -FlashInfer path over a 192-dim cache is exercised at M4/M6. - -Refs: vLLM ``models/deepseek_v2.py::DeepseekV2Attention`` (naive path) and -``rotary_embedding/deepseek_scaling_rope.py``. - -Run: pytest test/integration/test_kimi_mla.py -v -""" import pytest import torch import torch.nn.functional as F @@ -39,10 +20,6 @@ DEVICE = "cuda" -# -------------------------------------------------------------------------- -# References -# -------------------------------------------------------------------------- - def _ref_rmsnorm(x, weight, eps): x32 = x.float() x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) @@ -66,7 +43,6 @@ def _ref_yarn_rope(pos, q_pe, k_pe, rotary_dim, base, factor, max_pos, class _MockMLACache: - """Paged-cache stand-in: causal SDPA at 1/sqrt(head_dim).""" def __init__(self, head_dim: int): self.scale = head_dim ** -0.5 @@ -103,7 +79,6 @@ def _rope_kwargs(cfg): def _ref_mla(attn: KimiMLAAttention, cfg, h, pos, scale, boost): - """Independent MLA forward using weights extracted from ``attn``.""" T, H = h.shape[0], attn.num_heads Dnope, Drope, Dv, L = ( cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank) @@ -120,8 +95,7 @@ def _ref_mla(attn: KimiMLAAttention, cfg, h, pos, scale, boost): k_pe = k_pe.view(T, 1, Drope) rk = _rope_kwargs(cfg) q_pe, k_pe = _ref_yarn_rope(pos, q_pe, k_pe, **rk) - # M6 mitigation: q/k assembled at Dqk then zero-padded to padded_head_dim, v - # padded from Dv to padded_head_dim (see KimiMLAAttention.forward). + # Reference mirrors q/k/v zero-padding to padded_head_dim. pad = cfg.padded_head_dim q = F.pad(torch.cat([q_nope, q_pe], dim=-1), [0, pad - cfg.qk_head_dim]) * boost k = F.pad(torch.cat([k_nope, k_pe.expand(T, H, Drope)], dim=-1), [0, pad - cfg.qk_head_dim]) @@ -129,10 +103,6 @@ def _ref_mla(attn: KimiMLAAttention, cfg, h, pos, scale, boost): return q, k, v -# -------------------------------------------------------------------------- -# Tests -# -------------------------------------------------------------------------- - def test_yarn_rope_matches_reference(): torch.manual_seed(0) # mscale != mscale_all_dim so the cos/sin amplitude factor is non-trivial. diff --git a/test/integration/test_kimi_mla_absorb_forward.py b/test/integration/test_kimi_mla_absorb_forward.py index 45d3d5da0..659e517f6 100644 --- a/test/integration/test_kimi_mla_absorb_forward.py +++ b/test/integration/test_kimi_mla_absorb_forward.py @@ -1,20 +1,3 @@ -"""Phase-A GPU check: the REAL ``KimiMLAAttention.forward`` absorbed branch. - -The CPU gate (``test/modular/test_kimi_mla_absorb.py``) proves the absorption -algebra with pure-torch references. This test drives the actual wired forward — -``config.mla_absorb=True`` -> ``_forward_absorbed`` -> ``run_attention_mla`` — so -it needs a GPU (MLA RMSNorm uses a FlashInfer kernel). A ``_MockMLALatentCache`` -stands in for the Phase-B paged latent backend: its ``run_attention_mla`` does a -causal SDPA over ``[kv_c | k_pe]`` (value = ``kv_c``) at the DeepSeek scale, which -is exactly what the FlashInfer MLA kernel will compute. The real kernel is locked -to ckv=512/kpe=64 so it can't run at the reduced dims — that path is validated in -Phase B on real dims. - -Matching the independent DeepSeek MLA (materialized k_nope/v, no absorption) proves -the wired absorbed forward is numerically the naive path. - -Run: pytest test/integration/test_kimi_mla_absorb_forward.py -v -""" import pytest import torch import torch.nn.functional as F @@ -36,10 +19,6 @@ DEVICE = "cuda" -# -------------------------------------------------------------------------- -# References (device-aware; mirror test_kimi_mla_paged.py). -# -------------------------------------------------------------------------- - def _ref_rmsnorm(x, weight, eps): x32 = x.float() x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) @@ -81,7 +60,6 @@ def _deepseek_scale(cfg): def _ref_deepseek_mla(attn, cfg, h, pos): - """Naive DeepSeek MLA (materialized k_nope/v; no absorption) — ground truth.""" t, heads = h.shape[0], attn.num_heads d_nope, d_rope, d_v, latent = ( cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank) @@ -103,7 +81,6 @@ def _ref_deepseek_mla(attn, cfg, h, pos): class _MockMLALatentCache: - """Phase-B latent-backend stand-in: causal SDPA over [kv_c|k_pe], value=kv_c.""" def __init__(self, sm_scale): self.sm_scale = sm_scale diff --git a/test/integration/test_kimi_mla_absorb_kernel.py b/test/integration/test_kimi_mla_absorb_kernel.py index ee803cc6d..a699b37ce 100644 --- a/test/integration/test_kimi_mla_absorb_kernel.py +++ b/test/integration/test_kimi_mla_absorb_kernel.py @@ -1,23 +1,3 @@ -"""Phase-B follow-up GPU check: the FlashInfer MLA **kernel** fast path. - -``MlaAbsorbCacheManager`` uses ``flashinfer.mla.BatchMLAPagedAttentionWrapper`` -(the dedicated ckv=512/kpe=64 kernel) instead of the SDPA gather loop whenever -``_mla_kernel_available`` says the dims + GPU support it (real Kimi dims on sm90). -This drives the manager at real dims so the kernel activates and asserts: - - kernel path == SDPA path == independent accumulate-everything reference - -across a multi-page prefill, a decode step, and a batched (multi-request) decode. -The kernel manager (``mla_ckv_dim=512``) and the SDPA manager (``mla_ckv_dim=None``, -which forces the fallback at the SAME dims) share identical random inputs, so any -divergence is the kernel's. It also asserts the probe declines reduced dims (the -kernel is dim-locked and crashes off-dim, so reduced configs MUST use SDPA). - -Requires a Hopper (sm90) GPU — off sm90 the probe declines and the "kernel" -manager silently uses SDPA, making the comparison a tautology, so we skip. - -Run: pytest test/integration/test_kimi_mla_absorb_kernel.py -v -""" import pytest import torch @@ -41,14 +21,9 @@ ) DEVICE = torch.device("cuda") -# Real Kimi latent dims — the only dims the MLA kernel supports. L, DROPE = 512, 64 -# -------------------------------------------------------------------------- -# Real paged latent cache manager. ``ckv_dim=512`` -> kernel; None -> SDPA. -# -------------------------------------------------------------------------- - def _make_cm(softmax_scale, ckv_dim, request_ids, page_size=4, max_num_pages=64): latent_width = L + DROPE kv_cache = torch.zeros( @@ -87,30 +62,22 @@ def _make_cm(softmax_scale, ckv_dim, request_ids, page_size=4, max_num_pages=64) cm.set_layer_idx(0) return cm, alloc - -# -------------------------------------------------------------------------- -# Independent reference: accumulate every (kv_c, k_pe), causal SDPA (no paging). -# -------------------------------------------------------------------------- - def _ref_mla_step(q_nope_new, q_pe_new, kv_c_all, k_pe_all, scale): - """Intended MLA output for the NEW query tokens. ``kv_c_all``/``k_pe_all`` are - EVERY latent cached so far (including this step's); query j sits at absolute - position old_len+j and attends cached 0..old_len+j; value = the kv_c part.""" sl, _H, _L = q_nope_new.shape total = kv_c_all.shape[0] old_len = total - sl - query = torch.cat([q_nope_new, q_pe_new], dim=-1) # [sl,H,L+Drope] + query = torch.cat([q_nope_new, q_pe_new], dim=-1) key = torch.cat([kv_c_all.squeeze(1), k_pe_all.squeeze(1)], dim=-1) # [total,L+Drope] - value = kv_c_all.squeeze(1) # [total,L] - qt = query.transpose(0, 1).float() # [H,sl,L+Drope] + value = kv_c_all.squeeze(1) + qt = query.transpose(0, 1).float() scores = torch.einsum("hqd,kd->hqk", qt, key.float()) * scale q_pos = old_len + torch.arange(sl, device=DEVICE) k_pos = torch.arange(total, device=DEVICE) mask = torch.where(k_pos[None, :] <= q_pos[:, None], 0.0, torch.tensor(float("-inf"), device=DEVICE)) attn = (scores + mask).softmax(-1) - out = torch.einsum("hqk,kd->hqd", attn, value.float()) # [H,sl,L] - return out.transpose(0, 1).to(q_nope_new.dtype) # [sl,H,L] + out = torch.einsum("hqk,kd->hqd", attn, value.float()) + return out.transpose(0, 1).to(q_nope_new.dtype) def _rand_step(sl, H, dtype=torch.bfloat16): @@ -123,8 +90,6 @@ def _rand_step(sl, H, dtype=torch.bfloat16): def _run_single_req(cm, alloc, T, H, scale): - """Prefill T tokens (multi-page) + one decode step through ``cm``; return the - two outputs and the reference for each.""" q_nope, q_pe, kv_c, k_pe = _rand_step(T, H) cm.plan_attention(seq_lens=[T], is_causal=True, dtype=torch.bfloat16) with torch.no_grad(): @@ -146,12 +111,7 @@ def _run_single_req(cm, alloc, T, H, scale): def test_kernel_matches_sdpa_and_reference(): - """Real dims (L=512, Drope=64, H=2): kernel == SDPA == reference, prefill+decode. - - Same seed/inputs into a kernel manager (``mla_ckv_dim=512``) and an SDPA - manager (``mla_ckv_dim=None``) at identical dims — both must match the - independent reference and each other.""" - H, T, ps = 2, 6, 4 # T=6 over page_size-4 = 2 pages + H, T, ps = 2, 6, 4 scale = (L + DROPE) ** -0.5 * 1.3 assert _mla_kernel_available(L, DROPE, 9) is True @@ -162,45 +122,34 @@ def test_kernel_matches_sdpa_and_reference(): finally: alloc_k.cleanup() - torch.manual_seed(0) # identical inputs + torch.manual_seed(0) cm_s, alloc_s = _make_cm(scale, ckv_dim=None, request_ids=["r0"], page_size=ps) try: (sp, _refp), (sd, _refd) = _run_single_req(cm_s, alloc_s, T, H, scale) finally: alloc_s.cleanup() - # Sanity: the kernel manager actually took the kernel path (ps.wrapper set), - # the SDPA manager did not. assert cm_k._plan_states["main"].wrapper is not None assert cm_s._plan_states["main"].wrapper is None assert kp.shape == (T, H, L) and kd.shape == (1, H, L) - # kernel == reference torch.testing.assert_close(kp, refp, rtol=2e-2, atol=2e-2) torch.testing.assert_close(kd, refd, rtol=2e-2, atol=2e-2) - # SDPA == reference torch.testing.assert_close(sp, refp, rtol=2e-2, atol=2e-2) torch.testing.assert_close(sd, refd, rtol=2e-2, atol=2e-2) - # kernel == SDPA torch.testing.assert_close(kp, sp, rtol=2e-2, atol=2e-2) torch.testing.assert_close(kd, sd, rtol=2e-2, atol=2e-2) def test_kernel_batched_decode(): - """Batched (multi-request, varying lengths) decode through the kernel path. - - Two requests are prefilled to different lengths, then a single decode step - runs both together; each request's decode output must match its own - accumulate-everything reference.""" H, ps = 2, 4 scale = (L + DROPE) ** -0.5 - lens = [5, 9] # req0 prefill 5 tokens, req1 prefill 9 (spans >1 page) + lens = [5, 9] rids = ["r0", "r1"] torch.manual_seed(1) cm, alloc = _make_cm(scale, ckv_dim=L, request_ids=rids, page_size=ps) try: - # ---- prefill both requests (concatenated queries) ---- pref = [_rand_step(sl, H) for sl in lens] q_nope = torch.cat([p[0] for p in pref], 0) q_pe = torch.cat([p[1] for p in pref], 0) @@ -213,7 +162,6 @@ def test_kernel_batched_decode(): cm.advance_seq_lens() assert cm._plan_states["main"].wrapper is not None - # ---- one decode step for both requests ---- dec = [_rand_step(1, H) for _ in lens] dq_nope = torch.cat([d[0] for d in dec], 0) dq_pe = torch.cat([d[1] for d in dec], 0) @@ -225,9 +173,8 @@ def test_kernel_batched_decode(): torch.cuda.synchronize() assert got.shape == (2, H, L) - # ---- per-request reference (its prefill latents + its decode latent) ---- for i in range(2): - kv_c_all = torch.cat([pref[i][2], dec[i][2]], 0) # [len_i+1,1,L] + kv_c_all = torch.cat([pref[i][2], dec[i][2]], 0) k_pe_all = torch.cat([pref[i][3], dec[i][3]], 0) ref = _ref_mla_step(dec[i][0], dec[i][1], kv_c_all, k_pe_all, scale) torch.testing.assert_close(got[i:i + 1], ref, rtol=2e-2, atol=2e-2) @@ -236,17 +183,6 @@ def test_kernel_batched_decode(): def test_mla_wrapper_cuda_graph_capture_replay(): - """The MLA kernel wrapper + latent scatter are CUDA-graph capturable. - - Captures ONE decode graph (``set_latent`` + ``run`` over static buffers) then - replays it across two decode steps — re-planning each step (which updates the - wrapper's static index/scatter buffers via ``.copy_()``, advancing kv_len + - scatter offset) and copying fresh queries/latents into the static input - buffers before replay. Each replay's output must match the independent - accumulate-everything reference for that step. This mirrors exactly what - ``CudaGraphRunner`` does across decode steps, proving the absorbed-decode - primitive captures + replays correctly with a moving page table. - """ from mstar.utils.flashinfer_utils import FlashInferMLAWrapper bs, H, ps, max_pages = 2, 2, 4, 64 @@ -262,9 +198,7 @@ def test_mla_wrapper_cuda_graph_capture_replay(): max_total_tokens=bs, device=DEVICE, use_cuda_graph=True, ) - # Per-request fixed page ranges + prefill lengths. prefill=6/9 with ps=4 keep - # the two following decode steps within the same 2/3 pages (kv_indices stable; - # only kv_len + the scatter offset advance) — the common decode case. + # Decode steps stay within fixed pages; only kv_len and scatter offsets advance. req_pages = [[0, 1], [2, 3, 4]] prefill = [6, 9] torch.manual_seed(7) @@ -282,7 +216,6 @@ def test_mla_wrapper_cuda_graph_capture_replay(): kv_indices = torch.tensor(req_pages[0] + req_pages[1], device=DEVICE, dtype=torch.int32) qo_indptr = torch.tensor([0, 1, 2], device=DEVICE, dtype=torch.int32) - # Static input buffers replay reads from. q_nope_s = torch.zeros(bs, H, L, device=DEVICE, dtype=dtype) q_pe_s = torch.zeros(bs, H, DROPE, device=DEVICE, dtype=dtype) latent_s = torch.zeros(bs, latent_w, device=DEVICE, dtype=dtype) @@ -299,7 +232,6 @@ def fill_inputs(seed): q_pe_s.copy_(torch.randn(bs, H, DROPE, device=DEVICE, dtype=dtype) * 0.1) latent_s.copy_(torch.randn(bs, latent_w, device=DEVICE, dtype=dtype) * 0.1) - # ---- warmup (side stream) then capture ONE decode graph ---- plan_step(1) fill_inputs(100) s = torch.cuda.Stream() @@ -315,19 +247,17 @@ def fill_inputs(seed): wrapper.set_latent(cache, latent_s) out_static = wrapper.run(q_nope_s, q_pe_s, cache[..., :L], cache[..., L:]) - # ---- replay two decode steps; each must match its accumulate reference ---- decode_hist = [[] for _ in range(bs)] # decode latents scattered so far, per req prev = None for step in (1, 2): - fill_inputs(step) # fresh queries + decode latents into static bufs - plan_step(step) # re-plan: advance kv_len + scatter offset + fill_inputs(step) + plan_step(step) g.replay() torch.cuda.synchronize() for r in range(bs): decode_hist[r].append(latent_s[r:r + 1].clone()) got = out_static.clone() - # reference: query attends [prefix_r ; all decode latents so far] ref = torch.empty(bs, H, L, device=DEVICE, dtype=dtype) for r in range(bs): all_lat = torch.cat([prefix[r]] + decode_hist[r], dim=0) # [prefill+step, w] @@ -337,15 +267,12 @@ def fill_inputs(seed): ) torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2) if prev is not None: - # Sanity: distinct inputs -> distinct outputs (values really flow - # through the static buffers, not baked into the captured graph). + # Distinct inputs prove replay reads static buffers, not baked values. assert not torch.allclose(got, prev, atol=1e-3) prev = got def test_probe_declines_reduced_dims(): - """The kernel is dim-locked (ckv=512/kpe=64) and crashes off-dim, so the probe - must decline reduced dims / pre-sm90 (→ SDPA), and accept only real dims on sm90.""" assert _mla_kernel_available(L, DROPE, 9) is True # real dims, Hopper assert _mla_kernel_available(32, 8, 9) is False # reduced dims assert _mla_kernel_available(L, DROPE, 8) is False # pre-sm90 diff --git a/test/integration/test_kimi_mla_absorb_marlin_merge.py b/test/integration/test_kimi_mla_absorb_marlin_merge.py index c699c91a3..f98da5815 100644 --- a/test/integration/test_kimi_mla_absorb_marlin_merge.py +++ b/test/integration/test_kimi_mla_absorb_marlin_merge.py @@ -1,21 +1,3 @@ -"""Merge check: the absorbed-MLA and Marlin-MoE post-load hooks COMPOSE. - -The MLA weight-absorption work and the Marlin W4A16 MoE work were developed in -separate silos; both hang a ``process_weights_after_loading`` finalizer off the -same generic walker (``mstar.model.components.quantization.process_weights_after_loading``, -called once in ``kimi_model.py::_create_submodule``). This test builds a tree -holding BOTH a weight-absorbed ``KimiMLAAttention`` and a Marlin-backed -``KimiSparseMoeBlock`` and runs that single walker pass, asserting it finalizes -BOTH: the attention's ``w_kc``/``w_vc``/``fused_qkv_a_proj`` and the MoE's Marlin -repack (+ freed packed params). This is the production-default combination -(``k27_code`` => ``mla_absorb=True`` + ``quant_kernel="auto"`` => Marlin on sm80+), -which neither silo validated together. A Marlin-MoE forward post-walker confirms -the block still runs in the combined module. (Reduced dims => the attention uses -the SDPA-over-latent path, not the dim-locked FlashInfer MLA kernel; the kernel + -Marlin combined forward at real dims is covered by the TP8 serve smoke.) - -Run: pytest test/integration/test_kimi_mla_absorb_marlin_merge.py -v -""" import pytest import torch from torch import nn @@ -35,7 +17,6 @@ def _quantize_stack(weight): - """Per-expert fake-quantize a bf16 stack to packed int32 + bf16 scales.""" E, N, K = weight.shape packed = torch.empty((E, N, K // PACK_FACTOR), dtype=torch.int32, device=DEVICE) scale = torch.empty((E, N, K // GROUP_SIZE), dtype=torch.bfloat16, device=DEVICE) @@ -49,8 +30,6 @@ def _quantize_stack(weight): def _build_marlin_moe(cfg): - """A tp=1 Marlin-legal KimiSparseMoeBlock with synthetic packed experts - (mirrors test_kimi_moe_marlin._build_block).""" from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock with torch.device("meta"): @@ -75,7 +54,6 @@ def _build_marlin_moe(cfg): def _build_absorbed_attn(cfg): - """A materialized weight-absorbed KimiMLAAttention (random weights on CUDA).""" from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention torch.manual_seed(0) @@ -83,8 +61,6 @@ def _build_absorbed_attn(cfg): def test_generic_walker_composes_absorbed_mla_and_marlin_moe(): - """One ``process_weights_after_loading`` pass finalizes BOTH the absorbed MLA - projections and the Marlin MoE repack — the merged production default.""" cfg = KimiK2Config.reduced_marlin() cfg.mla_absorb = True # absorbed attention + Marlin experts in one config @@ -94,25 +70,20 @@ def test_generic_walker_composes_absorbed_mla_and_marlin_moe(): root.add_module("attn", attn) root.add_module("moe", moe) - # --- pre-walker: neither finalizer has run --- assert attn.mla_absorb assert attn.w_kc is None and attn.w_vc is None and attn.fused_qkv_a_proj_weight is None assert moe.experts.gate_up_proj_packed.numel() > 0 assert not getattr(moe, "_use_marlin", False) - # --- one generic walker pass (the exact call kimi_model._create_submodule makes) --- process_weights_after_loading(root, torch.device(DEVICE)) - # --- post-walker: absorbed-MLA finalizer ran --- assert attn.w_kc is not None, "walker did not build the absorbed w_kc" assert attn.w_vc is not None, "walker did not build the absorbed w_vc" assert attn.fused_qkv_a_proj_weight is not None, "walker did not build fused_qkv_a_proj" - # --- post-walker: Marlin-MoE finalizer ran --- assert moe._use_marlin, "walker did not select the Marlin backend" assert moe.experts.gate_up_proj_packed.numel() == 0, "packed experts not freed after repack" - # --- the Marlin MoE forward still runs in the combined module --- x = torch.randn(4, cfg.hidden_size, device=DEVICE, dtype=torch.bfloat16) * 0.1 with torch.no_grad(): out = moe(x) diff --git a/test/integration/test_kimi_mla_absorb_paged.py b/test/integration/test_kimi_mla_absorb_paged.py index 6e83307e3..21b8e7f94 100644 --- a/test/integration/test_kimi_mla_absorb_paged.py +++ b/test/integration/test_kimi_mla_absorb_paged.py @@ -1,21 +1,3 @@ -"""Phase-B GPU check: the REAL paged compressed-latent MLA backend. - -Drives ``MlaAbsorbCacheManager.run_attention_mla`` over a genuine 4D latent -paged cache (real ``PagedAllocationManager`` + ``create_cache_manager``), with -SYNTHETIC random ``q_nope/q_pe/kv_c/k_pe`` — no ``KimiMLAAttention`` needed. The -manager writes ``cat([kv_c, k_pe])`` as one latent vector per token into the -paged cache at its (page, offset), then per request gathers its full cached -latent and runs a causal SDPA (query = ``cat([q_nope, q_pe])``, key = the full -latent, value = its first ``L`` dims) at ``kv_cache_config.softmax_scale``. - -The reference is INDEPENDENT of the paging machinery: it accumulates every -``(kv_c, k_pe)`` seen so far into contiguous tensors and runs the same causal -SDPA (no pages, no scatter/gather). Matching it across a multi-page prefill and a -following decode step proves the paged scatter/gather + causal mask + scale are -correct — at real Kimi latent dims (L=512, Drope=64) and reduced dims. - -Run: pytest test/integration/test_kimi_mla_absorb_paged.py -v -""" import pytest import torch @@ -39,16 +21,10 @@ DEVICE = torch.device("cuda") -# -------------------------------------------------------------------------- -# Real paged latent cache manager (mirrors test_kimi_mla_paged.py, but 4D). -# -------------------------------------------------------------------------- - def _make_latent_cache_manager( latent_width, dtype, softmax_scale, page_size=4, max_num_pages=64 ): - # 4D latent cache: [num_layers, max_pages, page_size, latent_width] - # (the shape KVCacheEngine.load_model allocates for attention_backend - # "mla_absorb"). Small page_size so a handful of tokens spans >1 page. + # 4D mla_absorb cache; small page_size forces multi-page coverage. kv_cache = torch.zeros( 2, max_num_pages, page_size, latent_width, dtype=dtype, device=DEVICE, @@ -82,18 +58,7 @@ def _make_latent_cache_manager( return cm, alloc -# -------------------------------------------------------------------------- -# Independent reference: accumulate all (kv_c, k_pe), causal SDPA (no paging). -# -------------------------------------------------------------------------- - def _ref_mla_step(q_nope_new, q_pe_new, kv_c_all, k_pe_all, scale): - """The intended MLA output for the NEW query tokens. - - q_nope_new [sl,H,L], q_pe_new [sl,H,Drope] are just this step's queries; - kv_c_all [total,1,L], k_pe_all [total,1,Drope] are EVERY latent cached so far - (including this step's). Query j sits at absolute position old_len+j and - attends to cached 0..old_len+j; value is the kv_c (first L) part. - """ sl, _H, L = q_nope_new.shape total = kv_c_all.shape[0] old_len = total - sl @@ -128,8 +93,6 @@ def _rand_step(sl, H, L, Drope, dtype): def _run_prefill_then_decode(L, Drope, H, T, page_size): - """Drive a multi-page prefill + a decode step through the real paged backend - and compare each to the independent accumulate-everything reference.""" torch.manual_seed(0) dtype = torch.bfloat16 # Arbitrary MLA-style scale (the backend must apply exactly this value). @@ -140,7 +103,6 @@ def _run_prefill_then_decode(L, Drope, H, T, page_size): cm.set_active_label("main") cm.set_layer_idx(0) - # ---- prefill: T tokens spanning more than one page ---- assert T > page_size, "T must span >1 page to exercise page boundaries" q_nope, q_pe, kv_c, k_pe = _rand_step(T, H, L, Drope, dtype) cm.plan_attention(seq_lens=[T], is_causal=True, dtype=dtype) @@ -152,10 +114,8 @@ def _run_prefill_then_decode(L, Drope, H, T, page_size): assert got_prefill.shape == (T, H, L) torch.testing.assert_close(got_prefill, ref_prefill, rtol=2e-2, atol=2e-2) - # Advance seq_len so the decode step sees the T cached tokens. cm.advance_seq_lens() - # ---- decode: 1 new token attends over all T+1 cached ---- q_nope1, q_pe1, kv_c1, k_pe1 = _rand_step(1, H, L, Drope, dtype) cm.plan_attention(seq_lens=[1], is_causal=True, dtype=dtype) with torch.no_grad(): @@ -172,11 +132,8 @@ def _run_prefill_then_decode(L, Drope, H, T, page_size): def test_paged_latent_mla_real_dims(): - """Real Kimi MLA latent dims: L=512, Drope=64, H=2. Prefill (6 tokens over a - page_size-4 cache = 2 pages) + a decode step.""" _run_prefill_then_decode(L=512, Drope=64, H=2, T=6, page_size=4) def test_paged_latent_mla_reduced_dims(): - """Reduced dims: L=32, Drope=8, H=4. Same multi-page prefill + decode step.""" _run_prefill_then_decode(L=32, Drope=8, H=4, T=6, page_size=4) diff --git a/test/integration/test_kimi_mla_absorb_serve.py b/test/integration/test_kimi_mla_absorb_serve.py index b4c2da9ac..d444b1c83 100644 --- a/test/integration/test_kimi_mla_absorb_serve.py +++ b/test/integration/test_kimi_mla_absorb_serve.py @@ -1,22 +1,3 @@ -"""Absorbed-MLA serve smoke: the FULL model through the real mla_absorb backend. - -This is the end-to-end gate for the (now-default) weight-absorbed path. The -absorbed *math* is validated by ``test_kimi_mla_absorb.py`` (CPU) and the wired -attention forward by ``test_kimi_mla_absorb_forward.py``; the SDPA-over-latent -backend by ``test_kimi_mla_absorb_paged.py``. This test closes the loop: a whole -reduced ``KimiForCausalLM`` (all layers) built via the real -``get_submodule`` (meta -> to_empty -> load_weights -> post-load walker that builds -w_kc/w_vc + fused_qkv_a_proj) is driven for a prefill + several decode steps over a -genuine ``MlaAbsorbCacheManager`` + engine-shaped 4D latent cache. - -Correctness tie: the SAME synthetic checkpoint is loaded into a NAIVE reference -model (``mla_absorb=False``) and run through the mock cache at the DeepSeek scale; -the absorbed serve's prefill logits must match it (weight absorption is numerically -identical to naive up to fp rounding). This pins the absorbed serve path to the -M4-golden naive reference at full-model scale. - -Run: pytest test/integration/test_kimi_mla_absorb_serve.py -v -""" import pytest import torch @@ -42,11 +23,6 @@ DEVICE = torch.device("cuda") - -# -------------------------------------------------------------------------- -# Synthetic checkpoint (same serialization as test_kimi_submodule.py). -# -------------------------------------------------------------------------- - def _fill_layer(layer, cfg): a = layer.self_attn for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): @@ -70,7 +46,6 @@ def _fill_layer(layer, cfg): def _build_reference(cfg): - """A naive-path (mla_absorb=False) model with random weights = the ground truth.""" model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=torch.bfloat16) model.model.embed_tokens.weight.data.normal_(0, 0.05) model.model.norm.weight.data.normal_(1.0, 0.02) @@ -119,11 +94,6 @@ def _hf_checkpoint(model, cfg): sd["lm_head.weight"] = model.lm_head.weight return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()} - -# -------------------------------------------------------------------------- -# Naive mock cache (DeepSeek-correct scale) for the reference logits. -# -------------------------------------------------------------------------- - def _sdpa_causal(q, k, v, scale): qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) t = q.shape[0] @@ -145,11 +115,6 @@ def advance_seq_lens(self, *_a, **_k): def run_attention(self, q, k, v): return _sdpa_causal(q, k, v, self.scale) - -# -------------------------------------------------------------------------- -# Real mla_absorb backend (4D latent cache), engine-shaped. -# -------------------------------------------------------------------------- - def _absorbed_softmax_scale(cfg): r = cfg.rope_scaling mscale = yarn_get_mscale(r["factor"], r.get("mscale_all_dim", 0.0)) @@ -204,21 +169,15 @@ def _step(submodule, cm, graph_walk, token_ids): return out["logits"][0] -# -------------------------------------------------------------------------- -# Test -# -------------------------------------------------------------------------- - def test_absorbed_serve_matches_naive_reference(tmp_path): from safetensors.torch import save_file torch.manual_seed(0) - # One synthetic checkpoint, built from a naive reference model. cfg_naive = KimiK2Config.reduced() # mla_absorb=False (reduced default) assert cfg_naive.mla_absorb is False ref = _build_reference(cfg_naive) save_file(_hf_checkpoint(ref, cfg_naive), str(tmp_path / "model.safetensors")) - # Reference logits: the naive model through the mock cache at the DeepSeek scale. T = 6 prompt = torch.randint(0, cfg_naive.vocab_size, (T,), device=DEVICE) pos = torch.arange(T, device=DEVICE) @@ -226,14 +185,11 @@ def test_absorbed_serve_matches_naive_reference(tmp_path): ref_hidden = ref.model(prompt, _MockMLACache(cfg_naive.padded_head_dim), pos) ref_logits = ref.lm_head(ref_hidden[-1:]) # (1, vocab) - # Absorbed serve: load the SAME checkpoint into an mla_absorb model via the real - # build path; the post-load walker builds w_kc/w_vc + fused_qkv_a_proj. cfg_absorb = KimiK2Config.reduced() cfg_absorb.mla_absorb = True model = _make_model(cfg_absorb, tmp_path) submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16) assert isinstance(submodule, KimiLLMSubmodule) - # Absorbed models DO carry derived buffers (w_kc/w_vc/fused, persistent=False). buf_names = {n for n, _ in submodule.language_model.named_buffers()} assert any("w_kc" in n for n in buf_names) and any("fused_qkv_a_proj" in n for n in buf_names) @@ -242,10 +198,8 @@ def test_absorbed_serve_matches_naive_reference(tmp_path): prefill_logits = _step(submodule, cm, "prefill", prompt) assert prefill_logits.shape == (1, cfg_absorb.vocab_size) assert torch.isfinite(prefill_logits).all() - # Absorbed == naive up to bf16 rounding through a 2-layer stack + real backend. torch.testing.assert_close(prefill_logits, ref_logits, rtol=5e-2, atol=5e-2) - # A few decode steps over the accumulating paged LATENT cache. next_token = prefill_logits.argmax(-1) generated = [int(next_token.item())] for _ in range(4): diff --git a/test/integration/test_kimi_mla_paged.py b/test/integration/test_kimi_mla_paged.py index cbbbef1e9..101f246f8 100644 --- a/test/integration/test_kimi_mla_paged.py +++ b/test/integration/test_kimi_mla_paged.py @@ -1,32 +1,3 @@ -"""M6 step 1: the real paged MLA path, end-to-end, at the DeepSeek scale. - -This is the test that finally validates ``KimiMLAAttention`` over mstar's REAL -paged ``FlashInferCacheManager`` (genuine ``PagedAllocationManager`` + KV cache), -not the MockCacheHandle SDPA stand-in the M3/M4/M5 goldens use. - -It closes the M4 FlashInfer-192 blocker. The naive MLA pads q/k (from -``qk_head_dim``) and v (from ``v_head_dim``) up to ``padded_head_dim`` — the -smallest FlashInfer-SM90-supported head_dim {64,128,256} >= ``qk_head_dim`` — so -the reduced ``qk_head_dim=24`` becomes 64 (real Kimi 192 -> 256). The Hopper -prefill kernel ``static_assert``s ``head_dim_vo in {64,128,256}``, so the raw 24 -(and 192) fail to JIT-build; 64 builds and runs. - -The correctness crux is the **softmax-scale compensation**. run_attention applies -a fixed ``1/sqrt(padded_head_dim)`` scale, but DeepSeek's intended softmax scale -is ``qk_head_dim**-0.5 * mscale**2``. The zero-pad dims contribute 0 to q·k, so -we fold ``boost = mscale**2 * sqrt(padded_head_dim / qk_head_dim)`` into q: - - scores = (q*boost)·k * padded_head_dim**-0.5 - = q·k * mscale**2 * sqrt(padded/qk) * padded**-0.5 - = q·k * mscale**2 * qk**-0.5 (the DeepSeek scale). - -The reference below is the **independent DeepSeek computation** — projections + -YARN RoPE + causal SDPA at ``qk_head_dim**-0.5 * mscale**2`` over the UNPADDED q/k -(Dqk) and v (Dv), then output slice + o_proj. Matching it proves the padded paged -run + scale compensation reproduce the intended result exactly. - -Run: pytest test/integration/test_kimi_mla_paged.py -v -""" import pytest import torch import torch.nn.functional as F @@ -55,10 +26,6 @@ DEVICE = torch.device("cuda") -# -------------------------------------------------------------------------- -# Real paged cache manager (mirrors test_kimi_flashinfer_attention.py). -# -------------------------------------------------------------------------- - def _make_real_cache_manager(num_heads, head_dim, dtype, page_size=128, max_num_pages=8): kv_cache = torch.zeros( 2, max_num_pages, 2, page_size, num_heads, head_dim, @@ -91,10 +58,6 @@ def _make_real_cache_manager(num_heads, head_dim, dtype, page_size=128, max_num_ return cm, alloc -# -------------------------------------------------------------------------- -# Independent DeepSeek reference (no pad; scale = qk_head_dim**-0.5 * mscale**2). -# -------------------------------------------------------------------------- - def _ref_rmsnorm(x, weight, eps): x32 = x.float() x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) @@ -127,7 +90,6 @@ def _sdpa_causal(q, k, v, scale): def _ref_deepseek_mla(attn: KimiMLAAttention, cfg, h, pos): - """The intended DeepSeek MLA output: NO padding, scale = qk**-0.5 * mscale**2.""" T, H = h.shape[0], attn.num_heads Dnope, Drope, Dv, L = ( cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank) @@ -166,13 +128,6 @@ def _build_attention(cfg, dtype): def test_paged_mla_matches_deepseek_sdpa(): - """KimiMLAAttention through the REAL paged FlashInferCacheManager (head_dim = - padded_head_dim = 64) == the independent DeepSeek MLA at qk**-0.5 * mscale**2. - - This validates both (a) the real paged path builds+runs at the padded head_dim - (the M4 FlashInfer-192 blocker mitigation), and (b) the scale compensation is - exactly right — the padded run reproduces the unpadded DeepSeek scale. - """ torch.manual_seed(0) cfg = KimiK2Config.reduced() assert cfg.qk_head_dim == 24 and cfg.padded_head_dim == 64 # the mitigation @@ -196,6 +151,5 @@ def test_paged_mla_matches_deepseek_sdpa(): expected = _ref_deepseek_mla(attn, cfg, h, pos) assert got.shape == (T, cfg.hidden_size) - # bf16 through the real FlashInfer kernel; the scale compensation is exact in - # exact arithmetic, so any residual is pure bf16 rounding. + # Any residual after exact scale compensation is bf16 FlashInfer rounding. torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2) diff --git a/test/integration/test_kimi_moe.py b/test/integration/test_kimi_moe.py index 5714ff00b..e067f68d9 100644 --- a/test/integration/test_kimi_moe.py +++ b/test/integration/test_kimi_moe.py @@ -1,19 +1,3 @@ -"""M2 golden tests for Kimi-K2.7 fine-grained MoE. - -Verifies the new DeepSeek-V3 MoE math against independent references: - - the group-limited sigmoid ``noaux_tc`` router (KimiMoEGate), - - the fused expert dispatch (reused fused-expert GEMM), and - - the full MoE block (routed + ungated shared expert). - -References are inlined (self-contained; no dependency on the local golden -harness) and cited to vLLM ``fused_moe/cpu_fused_moe.py::grouped_topk`` and -``models/deepseek_v2.py::DeepseekV2MoE``. - -GPU test: the fused expert GEMM (``fused_experts``) is CUDA/bf16-only, so the -block/dispatch tests run on ``cuda``; the suite skips without a GPU. - -Run: pytest test/integration/test_kimi_moe.py -v -""" import pytest import torch import torch.nn.functional as F @@ -30,16 +14,10 @@ DEVICE = "cuda" - -# -------------------------------------------------------------------------- -# Independent references (cited to vLLM cpu_fused_moe.py / deepseek_v2.py) -# -------------------------------------------------------------------------- - def _ref_grouped_topk( logits: torch.Tensor, bias: torch.Tensor, n_group, topk_group, top_k, norm_topk_prob, routed_scaling_factor, ): - """vLLM ``grouped_topk`` (sigmoid + noaux_tc) -> (weights, ids).""" scores = logits.float().sigmoid() T = scores.shape[0] original = scores @@ -63,7 +41,6 @@ def _ref_grouped_topk( def _ref_routed_experts(h, gate_up, down, weights, ids): - """Naive per-token top-k expert loop matching ``fused_experts`` semantics.""" T, H = h.shape inter = down.shape[-1] out = torch.zeros(T, H, dtype=h.dtype, device=h.device) @@ -82,17 +59,10 @@ def _ref_swiglu(x, gate_w, up_w, down_w): def _dense_combine(ids, weights, num_experts): - """Scatter (ids, weights) into a dense (T, E) vector for order-insensitive - comparison (topk with sorted=False returns experts in arbitrary order).""" dense = torch.zeros(ids.shape[0], num_experts, device=ids.device) dense.scatter_(1, ids, weights.float()) return dense - -# -------------------------------------------------------------------------- -# Router -# -------------------------------------------------------------------------- - def test_moe_gate_matches_reference(): torch.manual_seed(0) cfg = KimiK2Config.reduced() @@ -122,8 +92,6 @@ def test_moe_gate_matches_reference(): def test_moe_gate_group_limited_routing(): - """With n_group=2/topk_group=1, every selected expert must come from the - single kept group — the crux of group-limited routing.""" torch.manual_seed(1) n_experts, n_group, topk_group, top_k = 8, 2, 1, 2 experts_per_group = n_experts // n_group @@ -141,11 +109,6 @@ def test_moe_gate_group_limited_routing(): groups = ids // experts_per_group assert (groups == groups[:, :1]).all(), "experts crossed group boundary" - -# -------------------------------------------------------------------------- -# Fused expert dispatch (trivial fixed router) -# -------------------------------------------------------------------------- - def test_expert_dispatch_matches_naive(): torch.manual_seed(2) cfg = KimiK2Config.reduced() @@ -155,7 +118,6 @@ def test_expert_dispatch_matches_naive(): h = torch.randn(T, H, device=DEVICE, dtype=dtype) * 0.1 gate_up = torch.randn(E, 2 * I, H, device=DEVICE, dtype=dtype) * 0.05 down = torch.randn(E, H, I, device=DEVICE, dtype=dtype) * 0.05 - # Trivial fixed router: every token -> experts {0, 1}, fixed weights. ids = torch.tensor([[0, 1]] * T, device=DEVICE) weights = torch.full((T, 2), 0.5, device=DEVICE, dtype=dtype) @@ -163,11 +125,6 @@ def test_expert_dispatch_matches_naive(): expected = _ref_routed_experts(h, gate_up, down, weights, ids) torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2) - -# -------------------------------------------------------------------------- -# Full MoE block (routed + ungated shared) -# -------------------------------------------------------------------------- - def test_moe_block_matches_reference(): torch.manual_seed(3) cfg = KimiK2Config.reduced() @@ -185,8 +142,7 @@ def test_moe_block_matches_reference(): sh_up = torch.randn(shared_inter, H, device=DEVICE, dtype=dtype) * 0.05 sh_down = torch.randn(H, shared_inter, device=DEVICE, dtype=dtype) * 0.05 - # Keep the router in fp32 (deterministic selection); load fused expert + - # shared weights. + # Keep the router fp32 for deterministic selection. block.gate.weight.data = gate_w block.gate.e_score_correction_bias.data = bias block.experts.gate_up_proj.data.copy_(expert_gate_up) diff --git a/test/integration/test_kimi_moe_inkernel_dequant.py b/test/integration/test_kimi_moe_inkernel_dequant.py index eaa966410..062813481 100644 --- a/test/integration/test_kimi_moe_inkernel_dequant.py +++ b/test/integration/test_kimi_moe_inkernel_dequant.py @@ -1,27 +1,3 @@ -"""GPU kernel golden: W4A16 in-kernel INT4 dequant vs the bf16 fused-expert GEMM. - -The in-kernel dequant path ships a SEPARATE ``fused_moe_kernel_w4a16`` that keeps the routed -experts packed in VRAM and dequantizes each K tile in registers before the dot. -Its correctness invariant is exact: the nibble ``(q - 8) * scale`` cast to bf16 is -*the same value* the bf16 path feeds to ``tl.dot`` after a pre-dequant, and with -the same tile config the two accumulate in the same order — so the packed path -must match the bf16 path on the SAME dequantized weights to a tight tolerance. - -This is the cheapest level that catches a kernel bug (packed-K stride, nibble -shifter, group-scale index, the top-nibble sign case) without a full model: - - 1. random bf16 experts ``w1 (E, 2I, H)`` / ``w2 (E, H, I)``; - 2. ``fake_quantize_weight`` each expert to ``(packed, bf16 scale, deq_bf16)``, - with ``scale_dtype=bfloat16`` so the packed-param scale and the bf16-path - weight dequantize from the identical scale; - 3. assert ``fused_experts(x, w1_packed, w2_packed, w1_scale=, w2_scale=, ...)`` - == ``fused_experts(x, w1_deq, w2_deq)`` (the bf16 path). - -Includes an expert whose packing sets container bit 31 (top nibble >= 8), proving -the arithmetic-shift + ``& 0xF`` mask recovers it. - -Run: pytest test/integration/test_kimi_moe_inkernel_dequant.py -v -""" import pytest import torch @@ -39,12 +15,6 @@ def _quantize_stack(weight): - """Fake-quantize a stacked ``(E, N, K)`` weight, returning packed/scale/deq. - - Each expert is quantized independently (matching a per-Linear checkpoint); - the bf16 scale is what a real compressed-tensors checkpoint stores, so the - returned ``deq`` is bit-for-bit what the packed param dequantizes to. - """ E, N, K = weight.shape packed = torch.empty((E, N, K // PACK_FACTOR), dtype=torch.int32, device=DEVICE) scale = torch.empty((E, N, K // GROUP_SIZE), dtype=torch.bfloat16, device=DEVICE) @@ -71,16 +41,14 @@ def test_w4a16_matches_bf16_on_same_dequant(num_tokens): torch.manual_seed(0) E, H, I, top_k = 4, 128, 64, 2 - # Slightly wide init so per-group amax spans the full nibble range and some - # top nibbles land >= 8 (container bit 31 set) — the sign-mask path. + # Wide init exercises top-nibble sign masking. w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16) w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16) w1_packed, w1_scale, w1_deq = _quantize_stack(w1) w2_packed, w2_scale, w2_deq = _quantize_stack(w2) - # Guard: the packing really exercises the negative-container / top-nibble>=8 - # case (else the sign-extension mask would be untested). + # Guard that the negative-container path is actually covered. assert (w1_packed < 0).any(), "no int32 with bit 31 set — top-nibble path untested" top_nibbles = unpack_int32(w1_packed.cpu(), num_bits=4)[..., PACK_FACTOR - 1 :: PACK_FACTOR] assert (top_nibbles >= 8).any() @@ -100,8 +68,6 @@ def test_w4a16_matches_bf16_on_same_dequant(num_tokens): def test_w4a16_reduce_results_false_shape(): - """``reduce_results=False`` returns the per-slot (tokens, top_k, hidden) tensor - the TP path all-reduces before the top-k sum — exercise it on the packed path.""" from mstar.utils.fused_moe.runner import fused_experts torch.manual_seed(1) diff --git a/test/integration/test_kimi_moe_marlin.py b/test/integration/test_kimi_moe_marlin.py index f42540820..34ce2909f 100644 --- a/test/integration/test_kimi_moe_marlin.py +++ b/test/integration/test_kimi_moe_marlin.py @@ -1,16 +1,3 @@ -"""GPU golden for the Marlin backend wired into ``KimiSparseMoeBlock``. - -Where ``test_marlin_kernels.py`` exercises the kernels in isolation, this drives -the *block* wiring: ``process_weights_after_loading`` resolving the backend + -repacking the packed experts (and freeing the source packed params), the -``_use_marlin`` branch in ``forward`` passing fp32 combine weights, and the shared -expert / router riding alongside. The reference is the same block's bf16 path: -router + bf16 fused-expert GEMM on the dequantized experts + the (bf16) shared -expert. Only the routed-expert kernel differs, so a cosine floor + relative-L2 -bound is the correct gate (see ``test_marlin_kernels.py`` for why). - -Run: pytest test/integration/test_kimi_moe_marlin.py -v -""" import pytest import torch @@ -41,8 +28,6 @@ def _quantize_stack(weight): def _build_block(): - """A tp=1 Marlin-legal KimiSparseMoeBlock, materialized on CUDA with random - router/shared weights and synthetic quantized routed experts loaded packed.""" from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock from mstar.model.kimi_k2_7.config import KimiK2Config @@ -53,13 +38,11 @@ def _build_block(): block = block.to(torch.bfloat16) block.to_empty(device=DEVICE) - # Random router + shared-expert weights (float params only). for p in block.parameters(): if p.dtype.is_floating_point: with torch.no_grad(): p.copy_(torch.randn_like(p) * 0.1) - # Synthetic quantized routed experts; keep the bf16 dequant for the reference. E, H, I = cfg.n_routed_experts, cfg.hidden_size, cfg.moe_intermediate_size w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16) w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16) @@ -80,17 +63,14 @@ def test_marlin_block_matches_bf16_reference(): H = cfg.hidden_size x = (torch.randn(5, H, device=DEVICE) * 0.5).to(torch.bfloat16) - # bf16 reference (before the Marlin repack frees the packed params): router + - # bf16 fused experts on the dequant + the shared expert. + # Build the bf16 reference before Marlin repack frees packed params. with torch.no_grad(): topk_w, topk_ids = block.gate(x) routed_ref = fused_experts(x, w1_deq, w2_deq, topk_w.to(x.dtype), topk_ids) ref = (routed_ref + block.shared_expert(x)).view(x.shape) - # Resolve backend + repack to Marlin, then run the block forward. block.process_weights_after_loading(torch.device(DEVICE)) assert block._use_marlin, "reduced_marlin config should select the Marlin backend" - # Source packed params are freed after the repack. assert block.experts.gate_up_proj_packed.numel() == 0 with torch.no_grad(): @@ -106,9 +86,6 @@ def test_marlin_block_matches_bf16_reference(): def test_forced_marlin_raises_on_illegal_shapes(): - """``quant_kernel='marlin'`` must fail loudly when the shapes are Marlin-illegal - (rather than silently downgrading) — reduced_quantized_inkernel has - moe_intermediate_size=64, which violates Marlin's k%128 on the down GEMM.""" from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock from mstar.model.kimi_k2_7.config import KimiK2Config @@ -123,7 +100,6 @@ def test_forced_marlin_raises_on_illegal_shapes(): def test_triton_backend_still_selected_when_forced(): - """``quant_kernel='triton'`` keeps the packed Triton path (no Marlin repack).""" from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock from mstar.model.kimi_k2_7.config import KimiK2Config diff --git a/test/integration/test_kimi_quant_inkernel_weight_loading.py b/test/integration/test_kimi_quant_inkernel_weight_loading.py index 5c730c8ac..52c5ec4c8 100644 --- a/test/integration/test_kimi_quant_inkernel_weight_loading.py +++ b/test/integration/test_kimi_quant_inkernel_weight_loading.py @@ -1,33 +1,3 @@ -"""Golden: packed experts + in-kernel INT4 dequant for Kimi-K2.7. - -Where the dequant-on-load golden (``test_kimi_quant_weight_loading.py``) -dequantizes every quantized weight to bf16 on load, in-kernel dequant keeps the -ROUTED EXPERTS packed int32 in VRAM and dequantizes each tile inside -``fused_moe_kernel_w4a16``; MLA / dense-FFN / shared-expert weights still dequant -on load. This test loads BOTH models from the SAME synthetic compressed-tensors -checkpoint and pins the packed-expert behavior against the dequant-on-load model: - - 1. build a reference (bf16 experts) and fake-quantize every eligible weight in - place, serializing an HF compressed-tensors checkpoint (``weight_packed`` + - ``weight_scale`` for quantized weights, plain ``weight`` for the ignore set); - 2. load ``model_a`` with ``reduced_quantized()`` (dequant-on-load — experts dequant - to bf16) and ``model_b`` with ``reduced_quantized_inkernel()`` (packed experts — - experts stay packed) from that one checkpoint; - 3. assert (a) completeness — ``model_b``'s loaded set equals its - ``named_parameters()``, now carrying ``*_packed`` / ``*_scale`` and NOT the - bf16 fused expert params; (b) the packed params survived the whole-model - ``.to(bf16)`` cast as int32 (the downcast-exemption guard); (c) router bias - fp32, no stray buffers; (d) ``model_b``'s full forward matches ``model_a``'s - within a loose bf16 tolerance (dequant-on-load and in-kernel dequant differ - only in accumulation order — both dequant to the identical bf16 weights). - -A tp=2 packed-MoE-block simulation (``test_packed_moe_block_tp2_sim_matches_tp1``) -separately proves the packed per-rank weight_loaders' column/row slicing and the -``_dispatch_packed_experts`` all-reduce path (shard_inter=32 satisfies the pack / -group divisibility asserts). - -Run: pytest test/integration/test_kimi_quant_inkernel_weight_loading.py -v -""" import pytest import torch @@ -47,10 +17,6 @@ DTYPE = torch.bfloat16 -# -------------------------------------------------------------------------- -# Mock paged cache (causal SDPA at 1/sqrt(head_dim)) — same as the dequant-on-load golden. -# -------------------------------------------------------------------------- - def _sdpa_causal(q, k, v, scale): qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) T = q.shape[0] @@ -72,11 +38,6 @@ def advance_seq_lens(self, *_a, **_k): def run_attention(self, q, k, v): return _sdpa_causal(q, k, v, self.scale) - -# -------------------------------------------------------------------------- -# Random weight init + fake-quant serialization (mirrors the dequant-on-load golden). -# -------------------------------------------------------------------------- - def _fill_layer(layer, cfg): a = layer.self_attn for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): @@ -90,8 +51,7 @@ def _fill_layer(layer, cfg): mlp.gate.weight.data.normal_(0, 1) mlp.gate.e_score_correction_bias.data = torch.randn( cfg.n_routed_experts, device=DEVICE, dtype=torch.float32) - # Wide expert init so per-group amax spans the nibble range (exercises the - # top-nibble sign path in the kernel) — the dequant-on-load build has bf16 experts. + # Wide init exercises the top-nibble sign path. mlp.experts.gate_up_proj.data.normal_(0, 0.2) mlp.experts.down_proj.data.normal_(0, 0.2) mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05) @@ -102,7 +62,6 @@ def _fill_layer(layer, cfg): def _build_reference(cfg): - """dequant-on-load config (bf16 fused experts) — the source the checkpoint is cut from.""" model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=DTYPE) model.model.embed_tokens.weight.data.normal_(0, 0.05) model.model.norm.weight.data.normal_(1.0, 0.02) @@ -114,7 +73,6 @@ def _build_reference(cfg): def _keep_bf16(key): - """Weights a real compressed-tensors Kimi checkpoint leaves in bf16.""" if key == "lm_head.weight": return True return ( @@ -126,8 +84,6 @@ def _keep_bf16(key): def _emit(sd, key, view, quant_cfg): - """Quantize (and write the dequant back into ``view``) when eligible, else store - the plain bf16 tensor — 2-D + input dim divisible by group_size + not bf16-kept.""" eligible = ( not _keep_bf16(key) and view.dim() == 2 @@ -195,7 +151,6 @@ def _hf_quant_checkpoint(model, cfg, quant_cfg): def _build_loaded(cfg, checkpoint_dir): - """Production path: meta -> to(bf16) -> to_empty(cuda) -> load_weights.""" with torch.device("meta"): model = KimiForCausalLM(cfg) model = model.to(DTYPE) @@ -203,34 +158,26 @@ def _build_loaded(cfg, checkpoint_dir): loaded = driver_load_weights(model, checkpoint_dir, device=DEVICE) return model.eval(), loaded - -# -------------------------------------------------------------------------- -# Test 1 — full-model packed-expert load + forward vs dequant-on-load. -# -------------------------------------------------------------------------- - def test_inkernel_weight_loading_and_forward_vs_dequant_on_load(tmp_path): from safetensors.torch import save_file torch.manual_seed(0) - cfg_a = KimiK2Config.reduced_quantized() # dequant-on-load (bf16 experts) - cfg_b = KimiK2Config.reduced_quantized_inkernel() # in-kernel dequant (packed experts) + cfg_a = KimiK2Config.reduced_quantized() + cfg_b = KimiK2Config.reduced_quantized_inkernel() assert cfg_b.moe_in_kernel_dequant and cfg_b.quantization_config is not None ref = _build_reference(cfg_a) - assert not isinstance(ref.model.layers[0].mlp, KimiSparseMoeBlock) # dense - assert isinstance(ref.model.layers[1].mlp, KimiSparseMoeBlock) # MoE + assert not isinstance(ref.model.layers[0].mlp, KimiSparseMoeBlock) + assert isinstance(ref.model.layers[1].mlp, KimiSparseMoeBlock) sd = _hf_quant_checkpoint(ref, cfg_a, cfg_a.quantization_config) - # The routed experts really are packed in the checkpoint (else in-kernel dequant is a no-op). assert any(k.endswith("mlp.experts.0.gate_proj.weight_packed") for k in sd) assert any(k.endswith("mlp.experts.0.down_proj.weight_packed") for k in sd) save_file(sd, str(tmp_path / "model.safetensors")) - model_a, _ = _build_loaded(cfg_a, tmp_path) # dequant-on-load reference - model_b, loaded_b = _build_loaded(cfg_b, tmp_path) # packed experts under test + model_a, _ = _build_loaded(cfg_a, tmp_path) + model_b, loaded_b = _build_loaded(cfg_b, tmp_path) - # (a) completeness: packed-expert loaded set == its named_parameters — includes the - # packed/scale params and EXCLUDES the bf16 fused expert params. all_params_b = set(dict(model_b.named_parameters()).keys()) assert loaded_b == all_params_b, ( f"unloaded: {all_params_b - loaded_b}; spurious: {loaded_b - all_params_b}") @@ -241,20 +188,16 @@ def test_inkernel_weight_loading_and_forward_vs_dequant_on_load(tmp_path): assert moe_prefix + "gate_up_proj" not in all_params_b # bf16 fused param gone assert moe_prefix + "down_proj" not in all_params_b - # (b) downcast-exemption guard: the packed params survived meta -> to(bf16) -> - # to_empty as int32 (PyTorch .to(dtype) skips integer tensors); scales are bf16. + # PyTorch .to(dtype) skips integer tensors; packed params must stay int32. experts_b = model_b.model.layers[1].mlp.experts assert experts_b.gate_up_proj_packed.dtype == torch.int32 assert experts_b.down_proj_packed.dtype == torch.int32 assert experts_b.gate_up_proj_scale.dtype == DTYPE assert experts_b.down_proj_scale.dtype == DTYPE - # (c) router bias fp32; no stray buffers survived either load path. assert model_b.model.layers[1].mlp.gate.e_score_correction_bias.dtype == torch.float32 assert {n for n, _ in model_b.named_buffers()} == set() - # Sanity: the two models' shared (non-expert) params are bit-identical, so any - # forward difference is isolated to the routed-expert path (in-kernel dequant vs dequant-on-load). a_sd = dict(model_a.named_parameters()) b_sd = dict(model_b.named_parameters()) shared_keys = set(a_sd) & set(b_sd) @@ -262,8 +205,6 @@ def test_inkernel_weight_loading_and_forward_vs_dequant_on_load(tmp_path): for name in shared_keys: assert torch.equal(a_sd[name], b_sd[name]), f"shared param mismatch at {name}" - # (d) full forward: in-kernel dequant matches dequant-on-load within a loose bf16 - # tolerance (they dequant to identical bf16 weights; residual is accumulation order only). T = 8 ids = torch.randint(0, cfg_b.vocab_size, (T,), device=DEVICE) pos = torch.arange(T, device=DEVICE) @@ -273,15 +214,7 @@ def test_inkernel_weight_loading_and_forward_vs_dequant_on_load(tmp_path): assert got.shape == (T, cfg_b.vocab_size) torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2) - -# -------------------------------------------------------------------------- -# Test 2 — tp=2 packed-MoE-block simulation (packed loaders' TP slicing + -# _dispatch_packed_experts all-reduce path). Mirrors test_kimi_tp.py's block sim. -# -------------------------------------------------------------------------- - class _NoCommGroup(CommGroup): - """world_size=2 comm group whose collectives are LOCAL no-ops: each rank - returns its partial and the test sums the two to reconstruct the reduce.""" def __init__(self, rank: int) -> None: super().__init__(my_global_rank=rank, my_group_rank=rank, group_members=[0, 1]) @@ -294,8 +227,6 @@ def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: def _packed_source(cfg, seed): - """Full-size expert source: bf16 gate/up/down, each fake-quantized to a full - packed+scale pair the per-rank loader then slices.""" g = torch.Generator().manual_seed(seed) E, H, I = cfg.n_routed_experts, cfg.hidden_size, cfg.moe_intermediate_size sh = I * cfg.n_shared_experts @@ -304,7 +235,7 @@ def _packed_source(cfg, seed): def rn(*shape, std=0.05, mean=0.0): return torch.randn(*shape, generator=g) * std + mean - def quant_stack(w): # (E, N, K) -> full packed (E,N,K//pf), scale (E,N,K//gs) + def quant_stack(w): packs, scales = [], [] for e in range(w.shape[0]): pk, sc, _ = fake_quantize_weight( @@ -351,7 +282,7 @@ def _load_moe_packed(block, src): def test_packed_moe_block_tp2_sim_matches_tp1(): - cfg = KimiK2Config.reduced_quantized_inkernel() # shard_inter=32 at tp=2 (ok) + cfg = KimiK2Config.reduced_quantized_inkernel() src = _packed_source(cfg, seed=707) g = torch.Generator().manual_seed(808) h = (torch.randn(7, cfg.hidden_size, generator=g) * 0.1).to(DEVICE, DTYPE) @@ -364,13 +295,12 @@ def test_packed_moe_block_tp2_sim_matches_tp1(): partials = [] for rank in range(2): block = KimiSparseMoeBlock(cfg, _NoCommGroup(rank)).to(DEVICE, DTYPE) - # each rank holds only a 1/2 stripe of the fused (packed) intermediate assert block.experts.gate_up_proj_packed.shape[1] == full_inter // 2 assert block.experts.down_proj_packed.dtype == torch.int32 _load_moe_packed(block, src) partials.append(block(h)) - out_tp2 = partials[0] + partials[1] # intermediate-parallel reduce == sum of ranks + out_tp2 = partials[0] + partials[1] max_abs = (out_tp2 - out_ref).abs().max().item() assert max_abs < 5e-2, f"packed MoE tp2 vs tp1 max abs diff {max_abs}" torch.testing.assert_close(out_tp2, out_ref, rtol=2e-2, atol=2e-2) diff --git a/test/integration/test_kimi_quant_weight_loading.py b/test/integration/test_kimi_quant_weight_loading.py index e6b19db32..f0bad8058 100644 --- a/test/integration/test_kimi_quant_weight_loading.py +++ b/test/integration/test_kimi_quant_weight_loading.py @@ -1,31 +1,3 @@ -"""Golden: compressed-tensors dequant-on-load for Kimi-K2.7. - -The real 1T INT4 checkpoint is absent and a bf16 dequant of it would not fit, so -this validates the *parser + numerics* on a SYNTHETIC ``reduced_quantized()`` -model, exactly mirroring the bf16 ``test_kimi_weight_loading.py`` but with a -compressed-tensors quantized checkpoint: - - 1. build a ``KimiForCausalLM`` reference with random bf16 weights; - 2. **fake-quantize in place** every eligible weight (2-D, input dim divisible by - ``group_size``, and not a norm / router / embedding / lm_head) — writing the - dequantized bf16 back into the reference, so the reference now holds exactly - what a correct loader must reproduce — and serialize it as an HF - compressed-tensors checkpoint: ``.weight_packed`` (int32) + - ``.weight_scale`` (bf16) for quantized weights, plain ``.weight`` - for the rest (the ``ignore`` set + weights whose dim doesn't divide the - group). MLA / dense-FFN / shared-expert / routed-expert weights are all - quantized here — proving dequant-on-load covers plain linears *and* the fused experts; - 3. build a fresh model on ``meta`` -> ``to(bf16)`` -> ``to_empty(cuda)`` and load - the checkpoint via the standard driver. Because the config carries a - ``quantization_config``, ``load_kimi_hf_weights`` wraps the stream with the - dequant-on-load dequantizer *before* the remap + stacked rules (which are unchanged); - 4. assert (a) completeness, (b) every loaded param equals the fake-quantized - reference bit-for-bit (the loader reproduces the dequant exactly), (c) the - router bias stays fp32 and no stray buffers survive, and (d) a full forward on - the loaded model matches a forward on the reference model. - -Run: pytest test/integration/test_kimi_quant_weight_loading.py -v -""" import pytest import torch @@ -43,10 +15,6 @@ DEVICE = "cuda" -# -------------------------------------------------------------------------- -# Mock paged cache (causal SDPA at 1/sqrt(head_dim)) — same as test_kimi_forward. -# -------------------------------------------------------------------------- - def _sdpa_causal(q, k, v, scale): qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) T = q.shape[0] @@ -69,11 +37,6 @@ def advance_seq_lens(self, *_a, **_k): def run_attention(self, q, k, v): return _sdpa_causal(q, k, v, self.scale) - -# -------------------------------------------------------------------------- -# Random weight init (router bias kept fp32) — same as test_kimi_weight_loading. -# -------------------------------------------------------------------------- - def _fill_layer(layer, cfg): a = layer.self_attn for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): @@ -103,19 +66,11 @@ def _build_reference(cfg): model.lm_head.weight.data.normal_(0, 0.02) for layer in model.model.layers: _fill_layer(layer, cfg) - # Eval-only fixture; disable grad so the in-place dequant write-back in _emit - # (copying into leaf params) is allowed. + # Disable grad so in-place dequant write-back into leaf params is allowed. model.requires_grad_(False) return model.eval() - -# -------------------------------------------------------------------------- -# Fake-quant serialization: compressed-tensors keys for eligible weights, plain -# bf16 for the rest. Mutates the reference in place to hold the dequant. -# -------------------------------------------------------------------------- - def _keep_bf16(key): - """Weights a real compressed-tensors Kimi checkpoint leaves in bf16.""" if key == "lm_head.weight": return True return ( @@ -127,12 +82,6 @@ def _keep_bf16(key): def _emit(sd, key, view, quant_cfg): - """Add checkpoint entries for reference param ``view`` at ``key``. - - Quantizes (and writes the dequant back into ``view``) when eligible, else - stores the plain bf16 tensor. Eligibility mirrors a real checkpoint: 2-D, - input dim divisible by the group size, and not in the bf16-keep set. - """ eligible = ( not _keep_bf16(key) and view.dim() == 2 @@ -141,9 +90,7 @@ def _emit(sd, key, view, quant_cfg): if not eligible: sd[key] = view return - # Store the scale in bf16 (as a real compressed-tensors checkpoint does) and - # take the dequant from that same bf16 scale, so the reference holds exactly - # what the loader's bf16-scale dequant reconstructs. + # Match compressed-tensors checkpoints: bf16 scale drives the reference dequant. packed, scale, deq = fake_quantize_weight( view, num_bits=quant_cfg.num_bits, group_size=quant_cfg.group_size, symmetric=quant_cfg.symmetric, scale_dtype=torch.bfloat16, @@ -204,7 +151,6 @@ def _hf_quant_checkpoint(model, cfg, quant_cfg): def _build_loaded(cfg, checkpoint_dir): - """Production path: meta -> to(bf16) -> to_empty(cuda) -> load_weights.""" with torch.device("meta"): model = KimiForCausalLM(cfg) model = model.to(torch.bfloat16) @@ -213,52 +159,38 @@ def _build_loaded(cfg, checkpoint_dir): return model.eval(), loaded -# -------------------------------------------------------------------------- -# Test -# -------------------------------------------------------------------------- - def test_quant_weight_loading_roundtrip_and_forward(tmp_path): from safetensors.torch import save_file torch.manual_seed(0) - cfg = KimiK2Config.reduced_quantized() # group_size=32, INT4 symmetric + cfg = KimiK2Config.reduced_quantized() assert cfg.quantization_config is not None ref = _build_reference(cfg) - # The stack spans the dense->MoE transition (first_k_dense_replace=1). assert not isinstance(ref.model.layers[0].mlp, KimiSparseMoeBlock) assert isinstance(ref.model.layers[1].mlp, KimiSparseMoeBlock) sd = _hf_quant_checkpoint(ref, cfg, cfg.quantization_config) - # Guard: the checkpoint really is quantized (else this silently degrades to - # the bf16 test) — routed experts AND plain linears carry packed weights. assert any(k.endswith("mlp.experts.0.gate_proj.weight_packed") for k in sd) assert any(k.endswith("self_attn.o_proj.weight_packed") for k in sd) - # ... and the ignore-set weights stayed bf16 (plain .weight, no packed). assert "lm_head.weight" in sd and "lm_head.weight_packed" not in sd assert "model.embed_tokens.weight" in sd save_file(sd, str(tmp_path / "model.safetensors")) model, loaded = _build_loaded(cfg, tmp_path) - # (a) completeness: every param received exactly one tensor, and no quant - # sub-key leaked through as a spurious param. all_params = set(dict(model.named_parameters()).keys()) assert loaded == all_params, ( f"unloaded: {all_params - loaded}; spurious: {loaded - all_params}") - # (b) every loaded param equals the fake-quantized reference, bit for bit - # (dequant-on-load reproduces the same fp32 (q-bias)*scale -> bf16 dequant). ref_sd = dict(ref.named_parameters()) for name, param in model.named_parameters(): assert torch.equal(param, ref_sd[name]), f"mismatch at {name}" - # (c) router bias preserved fp32; no derived buffers survived the load path. bias = model.model.layers[1].mlp.gate.e_score_correction_bias assert bias.dtype == torch.float32 assert {n for n, _ in model.named_buffers()} == set() - # (d) full forward on the loaded model matches the reference model's forward. T = 8 ids = torch.randint(0, cfg.vocab_size, (T,), device=DEVICE) pos = torch.arange(T, device=DEVICE) diff --git a/test/integration/test_kimi_serve_e2e.py b/test/integration/test_kimi_serve_e2e.py index 91cfaf173..537dd867a 100644 --- a/test/integration/test_kimi_serve_e2e.py +++ b/test/integration/test_kimi_serve_e2e.py @@ -1,37 +1,3 @@ -"""Phase 2 (gap 4): drive the Kimi-K2.7 SERVING path as deep as possible -in-process, beyond the submodule-level gate. - -``mstar-serve`` is a multi-process stack (API server -> conductor -> N worker -processes -> KV_CACHE engine -> decode Loop -> tokens over ZMQ). A live serve is -the ultimate proof; see ``tools/kimi_goldens/repro/RUNBOOK.md`` for the exact -launch + request commands. This test is the committed, deterministic gate that -exercises the same *model* serve surface a live serve hits, minus the inter- -process transport, so it runs in CI-style isolation on one GPU with synthetic -weights: - - 1. **The real serve entry points via the real __init__.** We build the model - through ``KimiK2Model(config_variant="reduced", checkpoint_path=..., - tokenizer_mode="byte")`` — the exact path ``api_server/entrypoint.py`` takes - from a serving YAML's ``model_kwargs`` — then use ``process_prompt`` (byte - tokenizer), ``get_submodule`` (meta -> to_empty -> M5 load), and - ``postprocess``. None of these are touched by ``test_kimi_submodule.py``, - which bypasses ``__init__`` via ``object.__new__``. - - 2. **prefill + the decode Loop with the real Sampler and check_stop.** Over a - genuine ``FlashInferCacheManager`` we run prefill (``forward`` -> logits -> - ``Sampler.sample``) then several decode steps through ``forward_batched`` - (the engine's batched decode path, which samples inside the pass and returns - per-request ``new_token``), calling the submodule's ``check_stop`` each step - — the Loop's real stop mechanism. This produces actual generated token ids - and proves the decode loop terminates on ``max_tokens`` (the reduced model's - EOS id 163586 is unreachable in ``vocab_size=256``, so ``max_tokens`` is the - only stop — exactly what a live reduced serve relies on). - -Deterministic by construction (greedy / ``temperature=0``): the sequence is -asserted stable across two fresh-cache runs so the golden never flakes. - -Run: pytest test/integration/test_kimi_serve_e2e.py -v -""" import pytest import torch @@ -59,11 +25,6 @@ DEVICE = torch.device("cuda") -# -------------------------------------------------------------------------- -# Synthetic HF DeepSeek-V3 reduced checkpoint (un-fuse every fused param back -# to HF keys — identical serialization to test_kimi_submodule / M5 loader). -# -------------------------------------------------------------------------- - def _fill_layer(layer, cfg): a = layer.self_attn for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): @@ -146,10 +107,6 @@ def _write_checkpoint(tmp_path, seed=0): return cfg -# -------------------------------------------------------------------------- -# Real paged FlashInfer cache (head_dim = padded_head_dim = 64 for reduced). -# -------------------------------------------------------------------------- - def _make_real_cache_manager(cfg, dtype, page_size=128, max_num_pages=8): num_heads = cfg.num_attention_heads head_dim = cfg.padded_head_dim @@ -179,7 +136,6 @@ def _make_real_cache_manager(cfg, dtype, page_size=128, max_num_pages=8): def _greedy_sampler(cfg): - """A real Sampler configured for deterministic greedy decode of r0.""" sampler = Sampler(device=DEVICE) sampler.add_request("r0") sampler.set_config("r0", vocab_size=cfg.vocab_size, temperature=0.0, @@ -188,8 +144,6 @@ def _greedy_sampler(cfg): def _fwd_info(max_tokens, cfg): - """CurrentForwardPassInfo the submodule's check_stop reads (sampling_config / - max_tokens / dynamic_loop_iter_counts).""" return CurrentForwardPassInfo( request_id="r0", graph_walk="decode", requires_cfg=False, fwd_index=0, random_seed=0, max_tokens=max_tokens, @@ -199,13 +153,7 @@ def _fwd_info(max_tokens, cfg): ) -# -------------------------------------------------------------------------- -# The serve drive. -# -------------------------------------------------------------------------- - def _run_generation(model, submodule, cfg, prompt_ids, max_tokens): - """prefill (forward+Sampler) then the decode Loop (forward_batched + check_stop) - over a fresh paged cache. Returns (generated_token_ids, stopped_by_check_stop).""" cm, alloc = _make_real_cache_manager(cfg, torch.bfloat16) sampler = _greedy_sampler(cfg) engine_inputs = ModelInputsFromEngine( @@ -215,7 +163,6 @@ def _run_generation(model, submodule, cfg, prompt_ids, max_tokens): generated: list[int] = [] stopped = False try: - # --- prefill: forward -> last-token logits -> Sampler (first token) --- ar = submodule.prepare_inputs("prefill", None, {"text_inputs": [prompt_ids]}) packed = submodule.preprocess("prefill", engine_inputs, [ar]) with torch.no_grad(): @@ -225,7 +172,6 @@ def _run_generation(model, submodule, cfg, prompt_ids, max_tokens): next_token = sampler.sample(["r0"], logits).clone() # (1,) generated.append(int(next_token.item())) - # --- decode Loop: forward_batched samples inside the pass; check_stop --- for step in range(max_tokens + 4): # +slack; check_stop must break first ar = submodule.prepare_inputs("decode", None, {"text_inputs": [next_token]}) packed = submodule.preprocess("decode", engine_inputs, [ar]) @@ -248,25 +194,15 @@ def _run_generation(model, submodule, cfg, prompt_ids, max_tokens): return generated, stopped -# -------------------------------------------------------------------------- -# Tests -# -------------------------------------------------------------------------- - def test_serve_path_prefill_decode_loop(tmp_path): - """Full model serve surface: real __init__ (reduced/local/byte) -> process_prompt - -> get_submodule -> prefill + decode Loop (Sampler + check_stop) -> postprocess. - Proves the decode loop generates real tokens and terminates on max_tokens.""" cfg = _write_checkpoint(tmp_path, seed=0) - # The exact construction api_server/entrypoint.py performs from a serving - # YAML's model_kwargs (no HF tokenizer, no 1T weights). model = KimiK2Model( model_path_hf="", config_variant="reduced", checkpoint_path=str(tmp_path), tokenizer_mode="byte", ) assert model.config.vocab_size == 256 - # Byte tokenizer: prompt text -> token ids in [0, 256). prompt_tensors = model.process_prompt("hello kimi", ["text"], ["text"]) prompt_ids = prompt_tensors["text_inputs"][0].to(DEVICE) assert prompt_ids.tolist() == list("hello kimi".encode("utf-8")) @@ -274,29 +210,22 @@ def test_serve_path_prefill_decode_loop(tmp_path): submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16) assert isinstance(submodule, KimiLLMSubmodule) - # M6 buffer audit still holds through the serve build path. assert list(submodule.language_model.named_buffers()) == [] MAX_TOKENS = 6 generated, stopped = _run_generation(model, submodule, cfg, prompt_ids, MAX_TOKENS) - # The decode loop stopped via check_stop (max_tokens), NOT the safety slack. assert stopped, "decode loop did not terminate via check_stop" - # 1 prefill token + exactly MAX_TOKENS decode tokens (check_stop fires when - # decode_loop count+1 >= max_tokens, i.e. after the MAX_TOKENS-th decode step). + # check_stop fires after exactly MAX_TOKENS decode steps. assert len(generated) == 1 + MAX_TOKENS, generated assert all(0 <= t < cfg.vocab_size for t in generated), generated - # postprocess decodes the generated ids back to bytes (client-facing output). out_bytes = model.postprocess(torch.tensor(generated), "text") assert isinstance(out_bytes, bytes) assert len(out_bytes) == len(generated) def test_serve_path_is_deterministic(tmp_path): - """Greedy decode over a fresh cache is bit-stable: two runs of the whole - prefill+decode serve drive yield identical token sequences. Guards the golden - against the flakiness that has bitten this project before.""" cfg = _write_checkpoint(tmp_path, seed=1) model = KimiK2Model( model_path_hf="", config_variant="reduced", diff --git a/test/integration/test_kimi_submodule.py b/test/integration/test_kimi_submodule.py index 19f2cf999..0f6fdec0d 100644 --- a/test/integration/test_kimi_submodule.py +++ b/test/integration/test_kimi_submodule.py @@ -1,32 +1,3 @@ -"""M6 step 6: submodule-level end-to-end through the REAL paged cache. - -This is the M6 correctness gate for serving: it exercises the whole -``KimiK2Model.get_submodule`` -> ``KimiLLMSubmodule`` -> real -``FlashInferCacheManager`` path on the reduced config with synthetic weights. - -Two things are validated: - -1. **The real build path.** ``get_submodule`` constructs ``KimiForCausalLM`` on - the meta device, casts to bf16 on meta, ``to_empty(cuda)``, and runs the M5 HF - loader — the production ``meta -> to_empty -> load_weights`` path (where the M5 - rope-buffer bug would have bitten). We assert the loaded model carries ZERO - buffers (M6 buffer audit) so no derived tensor survives as garbage. - -2. **Serving lifecycle over the paged MLA.** We drive - ``prepare_inputs -> preprocess -> forward`` through a genuine - ``FlashInferCacheManager`` (head_dim = padded_head_dim = 64) for a prefill plus - several decode steps, asserting sane token generation. The prefill logits are - checked against a mock-cache forward of the SAME loaded model at the - DeepSeek-correct scale, tying the paged serving path to the validated - MockCacheHandle goldens. - -``mstar-serve`` full-stack e2e (conductor + worker processes + SHM ports + CUDA -graph capture) is NOT run here — that infra isn't stood up in this environment. -This submodule-level test is the required correctness gate; see the M6 notes in -kimi-port-plan for the serve status. - -Run: pytest test/integration/test_kimi_submodule.py -v -""" import pytest import torch @@ -51,12 +22,6 @@ DEVICE = torch.device("cuda") - -# -------------------------------------------------------------------------- -# Synthetic HF DeepSeek-V3 checkpoint (same serialization as -# test_kimi_weight_loading — un-fuse every fused param back to HF keys). -# -------------------------------------------------------------------------- - def _fill_layer(layer, cfg): a = layer.self_attn for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): @@ -128,11 +93,6 @@ def _hf_checkpoint(model, cfg): sd["lm_head.weight"] = model.lm_head.weight return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()} - -# -------------------------------------------------------------------------- -# Real paged cache + mock cache (DeepSeek-correct scale via padded_head_dim). -# -------------------------------------------------------------------------- - def _make_real_cache_manager(cfg, dtype, page_size=128, max_num_pages=8): num_heads = cfg.num_attention_heads head_dim = cfg.padded_head_dim @@ -184,11 +144,6 @@ def run_attention(self, q, k, v): def _make_model(cfg, checkpoint_dir) -> KimiK2Model: - """A KimiK2Model wired to the synthetic checkpoint without a tokenizer. - - object.__new__ skips __init__ (which would pull a tokenizer / the full config), - so we set only what get_submodule needs — mirroring the modular test builder. - """ model = object.__new__(KimiK2Model) model.config = cfg model.model_path_hf = str(checkpoint_dir) @@ -196,11 +151,6 @@ def _make_model(cfg, checkpoint_dir) -> KimiK2Model: model._submodule_cache = {} return model - -# -------------------------------------------------------------------------- -# Minimal engine-inputs + lifecycle driver. -# -------------------------------------------------------------------------- - def _engine_inputs(cm): return ModelInputsFromEngine( request_ids=["r0"], per_request_info={}, cache_manager=cm, @@ -208,7 +158,6 @@ def _engine_inputs(cm): def _step(submodule, cm, graph_walk, token_ids): - """Drive prepare_inputs -> preprocess -> forward for one packed request.""" engine_inputs = _engine_inputs(cm) ar_in = submodule.prepare_inputs( graph_walk=graph_walk, fwd_info=None, @@ -220,10 +169,6 @@ def _step(submodule, cm, graph_walk, token_ids): return out["logits"][0], packed # (1, vocab), packed dict -# -------------------------------------------------------------------------- -# Test -# -------------------------------------------------------------------------- - def test_submodule_prefill_decode_over_real_paged_cache(tmp_path): from safetensors.torch import save_file @@ -232,18 +177,15 @@ def test_submodule_prefill_decode_over_real_paged_cache(tmp_path): ref = _build_reference(cfg) save_file(_hf_checkpoint(ref, cfg), str(tmp_path / "model.safetensors")) - # --- the real build path: meta -> to(bf16) -> to_empty(cuda) -> load --- model = _make_model(cfg, tmp_path) submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16) assert isinstance(submodule, KimiLLMSubmodule) - # get_submodule caches the built submodule. assert model.get_submodule("LLM") is submodule - # M6 buffer audit: no derived tensor buffer survived the load path as garbage. + # Derived buffers must not survive meta -> to_empty as garbage. assert list(submodule.language_model.named_buffers()) == [] p = next(submodule.language_model.parameters()) assert p.device.type == "cuda" and p.dtype == torch.bfloat16 - # --- prefill over the real paged FlashInfer cache (head_dim = 64) --- T = 6 prompt = torch.randint(0, cfg.vocab_size, (T,), device=DEVICE) cm, alloc = _make_real_cache_manager(cfg, torch.bfloat16) @@ -252,10 +194,7 @@ def test_submodule_prefill_decode_over_real_paged_cache(tmp_path): assert prefill_logits.shape == (1, cfg.vocab_size) assert torch.isfinite(prefill_logits).all() - # Reference: the SAME loaded model through the mock cache at the - # DeepSeek-correct scale (padded_head_dim). Ties the paged serving path to - # the validated MockCacheHandle goldens. Loose bf16 tolerance (2-layer stack - # through the real FlashInfer kernel). + # Same loaded weights through the mock cache at the DeepSeek-correct scale. pos = torch.arange(T, device=DEVICE) with torch.no_grad(): ref_hidden = submodule.language_model.model( @@ -263,7 +202,6 @@ def test_submodule_prefill_decode_over_real_paged_cache(tmp_path): ref_logits = submodule.lm_head(ref_hidden[-1:]) torch.testing.assert_close(prefill_logits, ref_logits, rtol=5e-2, atol=5e-2) - # --- a few decode steps over the accumulating paged KV cache --- next_token = prefill_logits.argmax(-1) # (1,) generated = [int(next_token.item())] assert 0 <= generated[-1] < cfg.vocab_size @@ -278,14 +216,11 @@ def test_submodule_prefill_decode_over_real_paged_cache(tmp_path): finally: alloc.cleanup() - # Sane generation: right length, all valid ids. assert len(generated) == 5 assert all(0 <= t < cfg.vocab_size for t in generated) def test_submodule_paged_decode_is_deterministic(tmp_path): - """Same prompt + fresh cache -> identical first token (paged path is stable and - the load is reproducible). Cheap guard against nondeterministic KV writes.""" from safetensors.torch import save_file torch.manual_seed(1) diff --git a/test/integration/test_kimi_tp.py b/test/integration/test_kimi_tp.py index 8ad47a1f0..35e763012 100644 --- a/test/integration/test_kimi_tp.py +++ b/test/integration/test_kimi_tp.py @@ -1,43 +1,3 @@ -"""Phase-3 tensor-parallel goldens for Kimi-K2.7 (reduced config): tp=2 == tp=1. - -Proves the two TP subsystems added in Phase 3 are numerically correct on the -reduced config: - - * **MLA head-sharding** (``KimiMLAAttention``): the q/kv UP-projections shard - ColumnParallel and ``o_proj`` reduces RowParallel, so each rank materializes - only its ``num_attention_heads // tp_size`` local heads. - * **MoE intermediate-sharding** (``KimiSparseMoeBlock``): the router stays - replicated; each rank holds every expert but only a - ``moe_intermediate_size // tp_size`` stripe of the SwiGLU intermediate - (gate_up column-parallel / down row-parallel), all-reduced before the top-k - sum-reduce. The shared expert shards via its ``ParallelGatedMLP``. - -Two verification levels, both keyed to the SAME deterministic source weights and -loaded through the REAL per-rank ``weight_loader`` slicing (one weight path for -tp=1 and tp>1): - - 1. **In-process rank simulation** (``test_*_tp2_sim_matches_tp1``, always runs - on one GPU): for each of the 2 ranks, build the block with ``tp_size=2`` and - that rank's weight shard, run it with a LOCAL no-op all-reduce so each rank - returns only its partial, then SUM the two ranks' partials and assert it - equals the tp=1 result within bf16 tolerance. This rigorously validates the - shard math + weight-loading; the block outputs are pure row-parallel reduces - (no un-sharded residual inside the block), so summing partials reconstructs - the reduce exactly. The live NCCL all-reduce itself rides the shared, - already-proven comm path (same primitive Orpheus tp2 uses). - - 2. **Real multi-process NCCL** (``test_tp2_nccl_matches_tp1``, runs only when - >= 2 CUDA devices are visible): spawn 2 ranks with a real NCCL comm group, - run attention + MoE + a full decoder layer with the REAL all-reduce, and - compare each rank's full output to the tp=1 reference. This exercises the - actual collective and the decoder layer's residual wiring (which the - partial-sum simulation cannot cover on its own). - -Determinism: all weights/inputs come from a fixed-seed CPU generator, so the max -abs diffs are stable run-to-run (no flakiness). - -Run: pytest test/integration/test_kimi_tp.py -v -""" from __future__ import annotations import os @@ -62,11 +22,6 @@ TP = 2 -# --------------------------------------------------------------------------- -# A world-size-2 comm group whose collectives are LOCAL no-ops. Used for the -# in-process rank simulation: each rank computes only its partial and the test -# sums the two ranks' partials to reconstruct the row-parallel all-reduce. -# --------------------------------------------------------------------------- class _NoCommGroup(CommGroup): def __init__(self, rank: int) -> None: super().__init__(my_global_rank=rank, my_group_rank=rank, group_members=[0, 1]) @@ -79,9 +34,6 @@ def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: class _MockMLACache: - """Paged-cache stand-in: causal SDPA at 1/sqrt(head_dim) over whatever local - heads the rank hands it (attention is per-head independent, so a rank running - SDPA on its head slice yields exactly those heads' outputs).""" def __init__(self, head_dim: int) -> None: self.scale = head_dim ** -0.5 @@ -107,10 +59,6 @@ def run_attention(self, q, k, v): return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype) -# --------------------------------------------------------------------------- -# Deterministic full-size source weights (CPU generator -> device-independent, -# identical across ranks / processes) and the REAL per-rank load helpers. -# --------------------------------------------------------------------------- def _source_weights(cfg: KimiK2Config, seed: int) -> dict: g = torch.Generator().manual_seed(seed) H = cfg.num_attention_heads @@ -121,7 +69,6 @@ def rn(*shape, std=0.03, mean=0.0): return torch.randn(*shape, generator=g) * std + mean return { - # attention (replicated down-projs + norms, sharded up-projs + o_proj) "q_a": rn(cfg.q_lora_rank, cfg.hidden_size), "q_a_norm": rn(cfg.q_lora_rank, std=0.02, mean=1.0), "q_b": rn(H * cfg.qk_head_dim, cfg.q_lora_rank), @@ -129,7 +76,6 @@ def rn(*shape, std=0.03, mean=0.0): "kv_a_norm": rn(cfg.kv_lora_rank, std=0.02, mean=1.0), "kv_b": rn(H * (cfg.qk_nope_head_dim + cfg.v_head_dim), cfg.kv_lora_rank), "o": rn(cfg.hidden_size, H * cfg.v_head_dim), - # moe (replicated fp32 router + sharded experts/shared) "router_w": torch.randn(E, Hd, generator=g), "router_b": torch.randn(E, generator=g), "gate": rn(E, I, Hd, std=0.05), @@ -138,15 +84,12 @@ def rn(*shape, std=0.03, mean=0.0): "sh_gate": rn(sh, Hd, std=0.05), "sh_up": rn(sh, Hd, std=0.05), "sh_down": rn(Hd, sh, std=0.05), - # decoder-layer norms "in_ln": rn(Hd, std=0.02, mean=1.0), "post_ln": rn(Hd, std=0.02, mean=1.0), } def _load_attention(attn: KimiMLAAttention, src: dict) -> None: - """Load full source weights through the REAL Column/Row weight_loaders (which - slice this rank's head block) + direct copies for the replicated params.""" attn.q_a_proj.weight.data.copy_(src["q_a"].to(DEVICE, DTYPE)) attn.q_a_layernorm.weight.data.copy_(src["q_a_norm"].to(DEVICE, DTYPE)) attn.q_b_proj.weight.weight_loader(attn.q_b_proj.weight, src["q_b"].to(DEVICE, DTYPE)) @@ -157,8 +100,6 @@ def _load_attention(attn: KimiMLAAttention, src: dict) -> None: def _load_moe(block: KimiSparseMoeBlock, src: dict) -> None: - """Load through the REAL fused-expert weight_loaders (per-rank intermediate - slice) + replicated fp32 router + the shared expert's merged/row loaders.""" block.gate.weight.data = src["router_w"].to(DEVICE) # keep router fp32 block.gate.e_score_correction_bias.data = src["router_b"].to(DEVICE) gu, dp = block.experts.gate_up_proj, block.experts.down_proj @@ -187,10 +128,6 @@ def _inputs(cfg: KimiK2Config, num_tokens: int, seed: int): pos = torch.arange(num_tokens, device=DEVICE) return h, pos - -# --------------------------------------------------------------------------- -# Level 1 — in-process rank simulation (single GPU, deterministic, always runs) -# --------------------------------------------------------------------------- def test_mla_attention_tp2_sim_matches_tp1(): cfg = KimiK2Config.reduced() src = _source_weights(cfg, seed=101) @@ -198,17 +135,17 @@ def test_mla_attention_tp2_sim_matches_tp1(): ref = KimiMLAAttention(cfg, CommGroup.trivial()).to(DEVICE, DTYPE) _load_attention(ref, src) - assert ref.num_heads == cfg.num_attention_heads # tp=1 sees all heads + assert ref.num_heads == cfg.num_attention_heads out_ref = ref(h, _MockMLACache(cfg.padded_head_dim), pos) partials = [] for rank in range(TP): attn = KimiMLAAttention(cfg, _NoCommGroup(rank)).to(DEVICE, DTYPE) - assert attn.num_heads == cfg.num_attention_heads // TP # rank sees local heads + assert attn.num_heads == cfg.num_attention_heads // TP _load_attention(attn, src) partials.append(attn(h, _MockMLACache(cfg.padded_head_dim), pos)) - out_tp2 = partials[0] + partials[1] # row-parallel o_proj reduce == sum of ranks + out_tp2 = partials[0] + partials[1] max_abs = (out_tp2 - out_ref).abs().max().item() assert max_abs < 5e-2, f"MLA tp2 vs tp1 max abs diff {max_abs}" torch.testing.assert_close(out_tp2, out_ref, rtol=2e-2, atol=2e-2) @@ -227,21 +164,17 @@ def test_moe_block_tp2_sim_matches_tp1(): partials = [] for rank in range(TP): block = KimiSparseMoeBlock(cfg, _NoCommGroup(rank)).to(DEVICE, DTYPE) - # each rank holds only a 1/TP stripe of the fused intermediate assert block.experts.gate_up_proj.shape[1] == full_inter // TP _load_moe(block, src) partials.append(block(h)) - out_tp2 = partials[0] + partials[1] # intermediate-parallel reduce == sum of ranks + out_tp2 = partials[0] + partials[1] max_abs = (out_tp2 - out_ref).abs().max().item() assert max_abs < 5e-2, f"MoE tp2 vs tp1 max abs diff {max_abs}" torch.testing.assert_close(out_tp2, out_ref, rtol=2e-2, atol=2e-2) def test_tp2_sim_is_stable_across_repeats(): - """The simulated tp2==tp1 diffs must be bit-stable across repeats (no flaky - golden). Re-run the attention + MoE simulation 3x from the same seeds and - assert an identical max abs diff each time.""" cfg = KimiK2Config.reduced() def attn_diff(): @@ -260,10 +193,6 @@ def attn_diff(): diffs = [attn_diff() for _ in range(3)] assert diffs[0] == diffs[1] == diffs[2], f"unstable tp2 sim diffs: {diffs}" - -# --------------------------------------------------------------------------- -# Level 2 — real multi-process NCCL (runs only with >= 2 CUDA devices) -# --------------------------------------------------------------------------- def _free_port() -> int: s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) s.bind(("127.0.0.1", 0)) @@ -273,12 +202,9 @@ def _free_port() -> int: def _nccl_worker(rank: int, world_size: int, port: int, result_path: str) -> None: - """One TP rank: real NCCL comm group, run attention + MoE + decoder layer, and - compare each to a tp=1 reference built in-process from the same source weights. - Rank 0 writes the max abs diffs to ``result_path``.""" import torch.distributed as dist - os.environ.setdefault("NCCL_IB_DISABLE", "1") # coriander has no RDMA/IB + os.environ.setdefault("NCCL_IB_DISABLE", "1") torch.cuda.set_device(rank) dist.init_process_group( backend="nccl", @@ -292,23 +218,21 @@ def _nccl_worker(rank: int, world_size: int, port: int, result_path: str) -> Non h, pos = _inputs(cfg, num_tokens=6, seed=606) cg = CommGroup(my_global_rank=rank, my_group_rank=rank, group_members=[0, 1]) - cg.device_group = None # 2-rank world == default group + cg.device_group = None cg.initialized = True - # tp=2 blocks with the REAL all-reduce -> each rank produces the FULL output. attn = KimiMLAAttention(cfg, cg).to(DEVICE, DTYPE) _load_attention(attn, src) moe = KimiSparseMoeBlock(cfg, cg).to(DEVICE, DTYPE) _load_moe(moe, src) dec = KimiDecoderLayer(cfg, layer_idx=1, comm_group=cg).to(DEVICE, DTYPE) _load_decoder(dec, src) - assert isinstance(dec.mlp, KimiSparseMoeBlock) # layer_idx=1 is a MoE layer + assert isinstance(dec.mlp, KimiSparseMoeBlock) attn_tp2 = attn(h, _MockMLACache(cfg.padded_head_dim), pos) moe_tp2 = moe(h) dec_tp2 = dec(h, _MockMLACache(cfg.padded_head_dim), pos) - # tp=1 reference (trivial group, same source weights). attn_ref = KimiMLAAttention(cfg, CommGroup.trivial()).to(DEVICE, DTYPE) _load_attention(attn_ref, src) moe_ref = KimiSparseMoeBlock(cfg, CommGroup.trivial()).to(DEVICE, DTYPE) @@ -344,8 +268,6 @@ def test_tp2_nccl_matches_tp1(tmp_path): mp.spawn(_nccl_worker, args=(TP, port, result_path), nprocs=TP, join=True) diffs = torch.load(result_path) - # Real all-reduce -> each rank's full output must match the tp=1 reference - # within bf16 tolerance (attention + MoE + decoder-layer residual wiring). assert diffs["attn"] < 5e-2, diffs assert diffs["moe"] < 5e-2, diffs assert diffs["decoder"] < 5e-2, diffs diff --git a/test/integration/test_kimi_weight_loading.py b/test/integration/test_kimi_weight_loading.py index 08b2849a7..26eb631f4 100644 --- a/test/integration/test_kimi_weight_loading.py +++ b/test/integration/test_kimi_weight_loading.py @@ -1,29 +1,3 @@ -"""M5 weight-loading golden test for Kimi-K2.7 / DeepSeek-V3 (synthetic checkpoint). - -The real 1T checkpoint is absent and would not fit, so this validates the loader -on a SYNTHETIC ``KimiK2Config.reduced()`` model with random bf16 weights: - - 1. build a ``KimiForCausalLM`` reference, fill it with random weights (router - ``e_score_correction_bias`` kept fp32); - 2. serialize it to a temp dir as an **HF DeepSeek-V3 checkpoint** — the exact - inverse of the loader's remap: per-expert ``gate_up_proj`` un-fused back to - ``experts.{e}.{gate,up}_proj``, dense/shared merged gate/up un-fused, singular - ``shared_expert`` -> HF plural ``shared_experts`` — as ``model.safetensors``; - 3. build a *fresh* model on ``meta`` -> ``to(bf16)`` -> ``to_empty(cuda)`` (the - production path), then load the checkpoint via the standard - ``mstar.model.loader.load_weights(model, dir, device)`` driver, which invokes - ``KimiForCausalLM.load_weights`` -> the M5 stacked rules + name remap; - 4. assert (a) every param of the loaded model equals the reference source - (exact), with targeted fused-param slice checks proving the gate/up/down - fusion for BOTH a dense (layer 0) and a MoE (layer 1) layer, and - (b) a full forward on the loaded model matches a forward on the reference - model (same mock-cache path as ``test_kimi_forward.py``). - -Confirms MLA loads strictly by name (no q_a/kv_a fusion): the attention params -appear identically in checkpoint and module, and the round-trip is exact. - -Run: pytest test/integration/test_kimi_weight_loading.py -v -""" import pytest import torch @@ -40,10 +14,6 @@ DEVICE = "cuda" -# -------------------------------------------------------------------------- -# Mock paged cache (causal SDPA at 1/sqrt(head_dim)) — same as test_kimi_forward. -# -------------------------------------------------------------------------- - def _sdpa_causal(q, k, v, scale): qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v)) T = q.shape[0] @@ -66,11 +36,6 @@ def advance_seq_lens(self, *_a, **_k): def run_attention(self, q, k, v): return _sdpa_causal(q, k, v, self.scale) - -# -------------------------------------------------------------------------- -# Random weight init (router bias kept fp32) + HF-checkpoint serialization. -# -------------------------------------------------------------------------- - def _fill_layer(layer, cfg): a = layer.self_attn for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj): @@ -83,7 +48,6 @@ def _fill_layer(layer, cfg): if isinstance(mlp, KimiSparseMoeBlock): # In place so the router weight keeps the model dtype (bf16). mlp.gate.weight.data.normal_(0, 1) - # Router selection bias stays fp32 even in a bf16 model. mlp.gate.e_score_correction_bias.data = torch.randn( cfg.n_routed_experts, device=DEVICE, dtype=torch.float32) mlp.experts.gate_up_proj.data.normal_(0, 0.05) @@ -106,11 +70,6 @@ def _build_reference(cfg): def _hf_checkpoint(model, cfg): - """Serialize the reference model to HF DeepSeek-V3 keys (inverse of the loader). - - Un-fuses every fused param back to the per-projection / per-expert checkpoint - layout so the loader has real fusion work to do. - """ inter = cfg.intermediate_size moe_inter = cfg.moe_intermediate_size shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts @@ -119,7 +78,6 @@ def _hf_checkpoint(model, cfg): for i, layer in enumerate(m.layers): p = f"model.layers.{i}." a = layer.self_attn - # MLA — identity keys, no fusion. sd[p + "self_attn.q_a_proj.weight"] = a.q_a_proj.weight sd[p + "self_attn.q_a_layernorm.weight"] = a.q_a_layernorm.weight sd[p + "self_attn.q_b_proj.weight"] = a.q_b_proj.weight @@ -153,7 +111,6 @@ def _hf_checkpoint(model, cfg): def _build_loaded(cfg, checkpoint_dir): - """Production path: meta -> to(bf16) -> to_empty(cuda) -> load_weights.""" with torch.device("meta"): model = KimiForCausalLM(cfg) model = model.to(torch.bfloat16) @@ -162,54 +119,38 @@ def _build_loaded(cfg, checkpoint_dir): return model.eval(), loaded -# -------------------------------------------------------------------------- -# Test -# -------------------------------------------------------------------------- - def test_weight_loading_roundtrip_and_forward(tmp_path): from safetensors.torch import save_file torch.manual_seed(0) cfg = KimiK2Config.reduced() ref = _build_reference(cfg) - # The stack spans the dense->MoE transition (first_k_dense_replace=1). assert not isinstance(ref.model.layers[0].mlp, KimiSparseMoeBlock) assert isinstance(ref.model.layers[1].mlp, KimiSparseMoeBlock) save_file(_hf_checkpoint(ref, cfg), str(tmp_path / "model.safetensors")) model, loaded = _build_loaded(cfg, tmp_path) - # (a0) completeness: every param received exactly one tensor. all_params = set(dict(model.named_parameters()).keys()) assert loaded == all_params, ( f"unloaded: {all_params - loaded}; spurious: {loaded - all_params}") - # (a1) every loaded param equals the reference source, bit for bit. ref_sd = dict(ref.named_parameters()) for name, param in model.named_parameters(): assert torch.equal(param, ref_sd[name]), f"mismatch at {name}" - # (a2) router bias preserved fp32 even in a bf16 model. bias = model.model.layers[1].mlp.gate.e_score_correction_bias assert bias.dtype == torch.float32 - # (a2b) regression guard (M6 buffer audit): NO derived tensor buffer survives - # meta -> to_empty as uninitialized garbage. The M6 audit of every Kimi - # submodule (attention/moe/decoder_layer/causal_lm/rope/language_model) found - # exactly one derived non-parameter tensor — the rope inv_freq — and it is - # computed lazily (M5 fix) rather than as an __init__ buffer, so the loaded - # model carries ZERO buffers. Any future __init__-computed buffer that is not - # in the checkpoint would fail this and silently corrupt the forward. + # Derived buffers must not survive meta -> to_empty as uninitialized memory. buffer_names = {n for n, _ in model.named_buffers()} assert buffer_names == set(), f"unexpected buffers survived the load path: {buffer_names}" - # (a3) targeted fusion checks — dense layer 0 (merged gate/up) ... inter = cfg.intermediate_size d_gup = model.model.layers[0].mlp.gate_up_proj.weight r_gup = ref.model.layers[0].mlp.gate_up_proj.weight assert torch.equal(d_gup[:inter], r_gup[:inter]) # gate half assert torch.equal(d_gup[inter:], r_gup[inter:]) # up half - # ... and MoE layer 1 (per-expert w13 gate|up + w2 down). mi = cfg.moe_intermediate_size l_gup = model.model.layers[1].mlp.experts.gate_up_proj r_egup = ref.model.layers[1].mlp.experts.gate_up_proj @@ -219,12 +160,6 @@ def test_weight_loading_roundtrip_and_forward(tmp_path): assert torch.equal(l_gup[e, mi:], r_egup[e, mi:]) # up:e assert torch.equal(model.model.layers[1].mlp.experts.down_proj, r_edwn) - # (b) full forward on the loaded model matches the reference model's forward. - # With bit-identical params (a1) AND a correctly-initialized rope (a2b), and - # since these kernels are per-instance deterministic (a repeated forward is - # bit-reproducible), the two forwards are bit-identical. The tiny bound below - # is off any tolerance boundary by ~3 orders of magnitude (measured diff is - # exactly 0.0 across runs) while still catching a gross mis-load (O(0.1+)). T = 8 ids = torch.randint(0, cfg.vocab_size, (T,), device=DEVICE) pos = torch.arange(T, device=DEVICE) diff --git a/test/integration/test_marlin_kernels.py b/test/integration/test_marlin_kernels.py index 0e9614bf1..8b012d1c4 100644 --- a/test/integration/test_marlin_kernels.py +++ b/test/integration/test_marlin_kernels.py @@ -1,22 +1,4 @@ -"""GPU golden for the vendored Marlin W4A16 kernels (utils/marlin). - -Validates the kernel layer in isolation, below any model wiring: - - 1. the JIT extension builds and registers ``_mstar_marlin_C`` ops; - 2. ``gptq_marlin_repack`` runs and is deterministic (repack layout is stable); - 3. the full routed-expert path (``MarlinMoEMethod`` = per-expert Marlin repack + - ``fused_marlin_moe``) matches the bf16 fused-expert GEMM on the *same* - dequantized weights, AND the existing Triton W4A16 path on the *same* packed - weights. - -Marlin dequantizes the identical INT4 nibbles as the Triton path but accumulates -in a different tile/reduce order (and folds fp32 combine weights), so agreement is -close-but-not-bit-exact. The meaningful gate is a cosine-similarity floor plus a -relative-L2 bound — an elementwise ``atol`` would flag the handful of bf16-accumulate -outliers on large down-projection sums, not a kernel bug. - -Run: pytest test/integration/test_marlin_kernels.py -v -""" +"""GPU golden tests for the vendored Marlin W4A16 kernels.""" import pytest import torch @@ -33,7 +15,6 @@ def _quantize_stack(weight): - """Fake-quantize a stacked ``(E, N, K)`` weight -> packed/scale/deq (bf16 scale).""" E, N, K = weight.shape packed = torch.empty((E, N, K // PACK_FACTOR), dtype=torch.int32, device=DEVICE) scale = torch.empty((E, N, K // GROUP_SIZE), dtype=torch.bfloat16, device=DEVICE) @@ -115,8 +96,6 @@ def test_marlin_moe_matches_bf16_and_triton(num_tokens): def test_marlin_moe_reduce_results_false_shape(): - """``reduce_results=False`` returns the per-slot (tokens, top_k, hidden) tensor - the TP path all-reduces before folding — exercise it on the Marlin path.""" from mstar.model.components.quantization import MarlinMoEMethod torch.manual_seed(1) diff --git a/test/modular/test_kimi_k27_code_wiring.py b/test/modular/test_kimi_k27_code_wiring.py index 90c270b75..684b33085 100644 --- a/test/modular/test_kimi_k27_code_wiring.py +++ b/test/modular/test_kimi_k27_code_wiring.py @@ -1,23 +1,3 @@ -"""CPU wiring tests for the REAL ``moonshotai/Kimi-K2.7-Code`` text-only serve. - -No weights, no GPU — the golden gate for the K2.7-Code serve plumbing while the -595 GB checkpoint is still downloading. Three concerns: - - 1. :meth:`KimiK2Config.k27_code` builds the full-size 1T text config with packed - experts armed and — crucially — keeps the default ``beta_fast=32.0`` (the - K2.7-Code ``text_config`` value). Guards against clobbering the YaRN field. - 2. :func:`kimi_name_remapper` strips the multimodal ``language_model.`` prefix so - the DeepSeek-V3 text keys land on ``KimiForCausalLM``'s params, routes the - packed routed-expert sub-keys through the packed-expert stacked rules, and drops the - vision (``vision_tower.*`` / ``mm_projector.*``) + ``weight_shape`` keys. A bare - ``model.*`` key (no prefix) is left unchanged. - 3. ``_maybe_apply_checkpoint_quant_config`` reads the ``quantization_config`` - NESTED under ``text_config`` (the multimodal wrapper leaves the top-level - null), while still parsing a flat top-level block and staying ``None`` for a - plain-bf16 checkpoint. - -Run: pytest test/modular/test_kimi_k27_code_wiring.py -v -""" import json import torch @@ -31,23 +11,17 @@ ) from mstar.model.loader.base import _apply_stacked -# -------------------------------------------------------------------------- -# 1. Config: k27_code() == full 1T dims + packed experts, default beta_fast=32.0. -# -------------------------------------------------------------------------- def test_k27_code_config_full_dims_packed_and_beta_fast(): cfg = KimiK2Config.k27_code() - # Packed experts armed, quant config auto-read from the checkpoint (still None here). assert cfg.moe_in_kernel_dequant is True assert cfg.quantization_config is None - # K2.7-Code keeps beta_fast=32.0 (guard against clobbering the YaRN field). assert cfg.rope_scaling["beta_fast"] == 32.0 assert cfg.rope_scaling["factor"] == 64.0 assert cfg.rope_scaling["rope_type"] == "deepseek_yarn" - # Full 1T text dims, matching the real Kimi-K2.7-Code text_config. assert cfg.num_hidden_layers == 61 assert cfg.n_routed_experts == 384 assert cfg.hidden_size == 7168 @@ -59,19 +33,12 @@ def test_k27_code_config_full_dims_packed_and_beta_fast(): assert cfg.qk_rope_head_dim == 64 assert cfg.v_head_dim == 128 - # It really is the full-size default plus exactly the one flag (no dim drift). base = KimiK2Config() assert cfg.num_hidden_layers == base.num_hidden_layers assert cfg.rope_scaling == base.rope_scaling # NO beta_fast override assert base.moe_in_kernel_dequant is False and cfg.moe_in_kernel_dequant is True - -# -------------------------------------------------------------------------- -# 2. Remapper: language_model.* strip + packed-expert routing + drops. -# -------------------------------------------------------------------------- - def _route(name, stacked): - """Mirror the loader: name_remapper then stacked-shard routing.""" mapped = kimi_name_remapper(name) if mapped is None: return None, None @@ -79,13 +46,12 @@ def _route(name, stacked): def test_remapper_language_model_prefix_and_packed_experts(): - cfg = KimiK2Config.reduced_quantized_inkernel() # in-kernel dequant => packed expert params + cfg = KimiK2Config.reduced_quantized_inkernel() with torch.device("meta"): model = KimiForCausalLM(cfg) params = set(dict(model.named_parameters()).keys()) stacked = build_kimi_stacked_params(cfg.n_routed_experts, packed_experts=True) - # -- plain text keys: strip language_model., land on a real param ------------ assert kimi_name_remapper( "language_model.model.layers.0.self_attn.q_a_proj.weight" ) == "model.layers.0.self_attn.q_a_proj.weight" @@ -101,14 +67,12 @@ def test_remapper_language_model_prefix_and_packed_experts(): ): assert landed in params - # -- shared expert: plural -> singular --------------------------------------- shared = kimi_name_remapper( "language_model.model.layers.1.mlp.shared_experts.down_proj.weight" ) assert shared == "model.layers.1.mlp.shared_expert.down_proj.weight" assert shared in params - # -- packed routed expert: remap + stacked -> the FOUR packed params --------- gate_p, gate_sid = _route( "language_model.model.layers.1.mlp.experts.0.gate_proj.weight_packed", stacked ) @@ -130,7 +94,6 @@ def test_remapper_language_model_prefix_and_packed_experts(): assert down_sid == "down:0" assert down_p in params - # -- vision drop: identity remap, NOT a model param (base loader skips it) ---- for vkey in ( "vision_tower.encoder.blocks.0.wqkv.weight", "mm_projector.proj.0.weight", @@ -139,23 +102,16 @@ def test_remapper_language_model_prefix_and_packed_experts(): target, _ = _route(vkey, stacked) assert target not in params # dropped - # -- weight_shape drop: routes to no real param ------------------------------ ws_target, _ = _route( "language_model.model.layers.1.mlp.experts.0.gate_proj.weight_shape", stacked ) assert ws_target not in params - # -- a flat model.* key (no language_model. prefix) is untouched ------------- assert ( kimi_name_remapper("model.layers.0.self_attn.q_a_proj.weight") == "model.layers.0.self_attn.q_a_proj.weight" ) - -# -------------------------------------------------------------------------- -# 3. Nested-quant reader: text_config.quantization_config + flat + bf16. -# -------------------------------------------------------------------------- - _QUANT_BLOCK = { "format": "pack-quantized", "quant_method": "compressed-tensors", @@ -176,15 +132,13 @@ def test_remapper_language_model_prefix_and_packed_experts(): def _make_model_with_config_json(tmp_dir, config_dict): - """A KimiK2Model with a bf16 (quant=None) config and a written config.json.""" (tmp_dir / "config.json").write_text(json.dumps(config_dict)) model = object.__new__(KimiK2Model) - model.config = KimiK2Config() # quantization_config is None by default + model.config = KimiK2Config() return model def test_nested_quant_config_read(tmp_path): - # Nested under text_config, top-level absent — the real K2.7-Code layout. d = tmp_path / "nested" d.mkdir() model = _make_model_with_config_json( @@ -200,7 +154,6 @@ def test_nested_quant_config_read(tmp_path): def test_flat_quant_config_read_backward_compat(tmp_path): - # A flat top-level quantization_config block must still parse. d = tmp_path / "flat" d.mkdir() model = _make_model_with_config_json(d, {"quantization_config": _QUANT_BLOCK}) @@ -212,7 +165,6 @@ def test_flat_quant_config_read_backward_compat(tmp_path): def test_plain_bf16_config_stays_none(tmp_path): - # No quant block anywhere (nested or flat) — stays bf16 (None). d = tmp_path / "bf16" d.mkdir() model = _make_model_with_config_json(d, {"text_config": {"num_hidden_layers": 61}}) diff --git a/test/modular/test_kimi_mla_absorb.py b/test/modular/test_kimi_mla_absorb.py index ff8b98201..3b93aacaf 100644 --- a/test/modular/test_kimi_mla_absorb.py +++ b/test/modular/test_kimi_mla_absorb.py @@ -1,27 +1,3 @@ -"""Phase-A confirmation gate for weight-absorbed MLA (``config.mla_absorb``). - -These are CPU-only, GPU-free tests that prove the absorption *math* on the reduced -config, behind the default-off flag. They deliberately do NOT call -``KimiMLAAttention.forward`` (its RMSNorm uses a FlashInfer GPU kernel); instead -they use pure-torch references (mirroring ``test/integration/test_kimi_mla*.py``) -and exercise the real load-time build (``process_weights_after_loading``) + the -absorption algebra + the latent KV-cache config. - -What the absorbed path must satisfy: - * ``w_kc``/``w_vc`` split out of ``kv_b_proj`` reconstruct the naive per-head - ``k_nope``/``v`` from the latent (``test_absorb_reconstructs_kv_b_proj``); - * the absorbed forward math equals the canonical DeepSeek MLA output — the same - thing the naive path reproduces (``test_absorbed_math_matches_deepseek``); - * ``get_kv_cache_config`` reports the shrunk latent cache when the flag is on, - and is byte-identical to naive when off (``test_kv_cache_config_*``). - -The real FlashInfer MLA kernel / paged latent cache (Phase B) is out of scope here -(it is hard-locked to ckv=512/kpe=64, so it cannot run at the reduced dims); the -wired ``forward`` absorbed branch is exercised on GPU in -``test/integration/test_kimi_mla_absorb_forward.py``. - -Run: pytest test/modular/test_kimi_mla_absorb.py -v -""" import torch import torch.nn.functional as F @@ -35,10 +11,6 @@ from mstar.model.kimi_k2_7.config import KimiK2Config from mstar.model.kimi_k2_7.kimi_model import KimiK2Model -# -------------------------------------------------------------------------- -# Pure-torch references (CPU; copied from test_kimi_mla_paged.py so this stays -# self-contained and GPU-free). -# -------------------------------------------------------------------------- def _ref_rmsnorm(x, weight, eps): x32 = x.float() @@ -75,7 +47,6 @@ def _sdpa_causal(q, k, v, scale): def _q_and_latent(attn, cfg, h, pos): - """Shared Q + normed-latent + roped pe slices (the piece both refs need).""" t, heads = h.shape[0], attn.num_heads d_nope, d_rope, latent = cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.kv_lora_rank eps = cfg.rms_norm_eps @@ -97,7 +68,6 @@ def _deepseek_scale(cfg): def _ref_deepseek_mla(attn, cfg, h, pos): - """Canonical (naive) DeepSeek MLA: materialize k_nope/v, SDPA at Dqk, o_proj.""" t, heads = h.shape[0], attn.num_heads d_nope, d_rope, d_v = cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim q_nope, q_pe, kv_c, k_pe = _q_and_latent(attn, cfg, h, pos) @@ -110,7 +80,6 @@ def _ref_deepseek_mla(attn, cfg, h, pos): def _absorbed_mla(attn, cfg, h, pos): - """Weight-absorbed MLA: fold w_kc into q, MQA over the latent, fold w_vc into o.""" t, heads = h.shape[0], attn.num_heads d_rope, d_v, latent = cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank q_nope, q_pe, kv_c, k_pe = _q_and_latent(attn, cfg, h, pos) @@ -124,7 +93,6 @@ def _absorbed_mla(attn, cfg, h, pos): def _build_attention_cpu(seed=0): - """Reduced-config KimiMLAAttention on CPU (fp32), absorbed weights built.""" torch.manual_seed(seed) cfg = KimiK2Config.reduced() cfg.mla_absorb = True @@ -138,12 +106,7 @@ def _build_attention_cpu(seed=0): return attn, cfg -# -------------------------------------------------------------------------- -# Tests -# -------------------------------------------------------------------------- - def test_absorb_reconstructs_kv_b_proj(): - """w_kc/w_vc built from kv_b_proj reproduce naive per-head k_nope / v.""" attn, cfg = _build_attention_cpu(seed=1) heads, d_nope, d_v, latent = ( attn.num_heads, cfg.qk_nope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank) @@ -161,21 +124,16 @@ def test_absorb_reconstructs_kv_b_proj(): def test_fused_qkv_a_proj(): - """The fused latent down-projection buffer == cat(q_a_proj, kv_a_proj_with_mqa) - exactly (it is the one GEMM the absorbed forward runs before splitting).""" attn, cfg = _build_attention_cpu(seed=3) expected = torch.cat( [attn.q_a_proj.weight, attn.kv_a_proj_with_mqa.weight], dim=0) assert tuple(attn.fused_qkv_a_proj_weight.shape) == ( cfg.q_lora_rank + cfg.kv_lora_rank + cfg.qk_rope_head_dim, cfg.hidden_size) - # Byte-for-byte concat; no arithmetic, so exact equality (rtol/atol 0). torch.testing.assert_close( attn.fused_qkv_a_proj_weight, expected, rtol=0, atol=0) def test_absorbed_math_matches_deepseek(): - """The absorbed forward math == the canonical DeepSeek MLA output (the naive - invariant). This is the Phase-A algorithm gate.""" attn, cfg = _build_attention_cpu(seed=2) t = 7 h = torch.randn(t, cfg.hidden_size) * 0.1 @@ -185,8 +143,7 @@ def test_absorbed_math_matches_deepseek(): reference = _ref_deepseek_mla(attn, cfg, h, pos) assert absorbed.shape == (t, cfg.hidden_size) - # Pure fp32 algebra; the only difference is float op ordering (two bmms + - # latent SDPA vs materialized SDPA), so the residual is tiny. + # Pure fp32 algebra; residual comes only from op ordering. torch.testing.assert_close(absorbed, reference, rtol=1e-4, atol=1e-4) @@ -204,7 +161,6 @@ def test_kv_cache_config_absorbed_shrinks_latent(): assert kv.head_dim == cfg.kv_lora_rank + cfg.qk_rope_head_dim # 32 + 8 = 40 assert kv.num_qo_heads == cfg.num_attention_heads # 4 (q still sharded) assert kv.attention_backend == "mla_absorb" - # per-token cache shrink vs naive padded MHA: 2 * 4 * 64 = 512 -> 1 * 40 = 40 naive_elems = 2 * cfg.num_attention_heads * cfg.padded_head_dim absorbed_elems = kv.num_kv_heads * kv.head_dim assert naive_elems == 512 and absorbed_elems == 40 diff --git a/test/modular/test_kimi_model.py b/test/modular/test_kimi_model.py index 0cca57cc7..baef9eac9 100644 --- a/test/modular/test_kimi_model.py +++ b/test/modular/test_kimi_model.py @@ -1,11 +1,3 @@ -"""M0 scaffold tests for Kimi-K2.7 (text backbone). - -Dummy mode: the model is built via ``object.__new__`` (no tokenizer, no weights, -no GPU) and only the ``Model`` contract is exercised — the graph, engine types, -KV-cache dims, and the prefill→decode→done state machine. This validates the -serving plumbing in isolation before any MLA/MoE compute exists, exactly as -``docs/adding_models.rst`` prescribes for a new model. -""" import sys sys.path.insert(0, ".") @@ -44,13 +36,9 @@ def test_kimi_kv_cache_config_matches_reduced_mla_dims(): (kv,) = kv assert kv.num_layers == cfg.num_hidden_layers == 2 - # Naive/materialized MLA: KV heads == query heads. assert kv.num_kv_heads == cfg.num_attention_heads == 4 assert kv.num_qo_heads == cfg.num_attention_heads == 4 - # M6 FlashInfer-SM90 mitigation: q/k/v are zero-padded from qk_head_dim (24) - # up to the smallest supported head_dim {64,128,256} >= qk_head_dim, so the - # paged cache stores head_dim == padded_head_dim == 64 (not the raw 24, which - # the Hopper prefill kernel static_asserts against). + # FlashInfer-SM90 requires padded_head_dim, not raw qk_head_dim. assert cfg.qk_head_dim == cfg.qk_nope_head_dim + cfg.qk_rope_head_dim == 24 assert kv.head_dim == cfg.padded_head_dim == 64 assert kv.max_seq_len == cfg.max_position_embeddings @@ -98,9 +86,5 @@ def test_kimi_decode_completion_marks_done(): def test_kimi_get_submodule_is_dummy_mode(): model = _make_model() - # M6: get_submodule is the real meta->to_empty->load_weights build, but it - # returns None in dummy mode when no checkpoint is resolvable. _make_model sets - # no model_path_hf, so _resolve_checkpoint() -> None -> dummy mode, letting the - # modular graph tests run without a GPU or weights. assert getattr(model, "model_path_hf", None) is None assert model.get_submodule("LLM") is None diff --git a/test/modular/test_kimi_quant.py b/test/modular/test_kimi_quant.py index aed5f4df3..390a649db 100644 --- a/test/modular/test_kimi_quant.py +++ b/test/modular/test_kimi_quant.py @@ -1,23 +1,3 @@ -"""CPU unit tests for the Kimi-K2.7 compressed-tensors dequant utilities (dequant-on-load). - -These pin the *numerics and bit layout* of the dequant-on-load parser without a -GPU or checkpoint — the cheapest level that guards the correctness harness: - - 1. known-answer pack/unpack (the exact int32 bit layout, incl. the sign-bit - nibble), so a future refactor can't silently change the on-disk convention; - 2. pack/unpack round-trip on random nibbles; - 3. symmetric dequant math (offset-binary ``(nibble - bias) * scale``); - 4. ``fake_quantize_weight`` -> ``dequantize_weight`` exactness (the golden - harness relies on the loader reproducing the fake-quant result bit-for-bit); - 5. the streaming generator: quant components collapse to one bf16 ``*.weight``, - non-quant keys pass through, incomplete groups raise; - 6. ``CompressedTensorsQuantConfig.from_hf_config_dict`` parsing. - -The full weight-loading + forward golden (needs the fused-expert GEMM + RMSNorm) -lives in ``test/integration/test_kimi_quant_weight_loading.py``. - -Run: pytest test/modular/test_kimi_quant.py -v -""" import pytest import torch @@ -30,21 +10,15 @@ unpack_int32, ) -# -------------------------------------------------------------------------- -# 1. Known-answer pack/unpack — pins the int32 bit layout. -# -------------------------------------------------------------------------- def test_pack_known_answer(): - # Eight INT4 nibbles 0..7 along the last axis pack low-order-first: - # sum(j << 4*j for j in 0..7) == 0x76543210. nibbles = torch.arange(8, dtype=torch.int64).reshape(1, 8) packed = pack_int32(nibbles, num_bits=4) assert packed.dtype == torch.int32 assert packed.shape == (1, 1) assert packed.item() == 0x76543210 - # A top nibble >= 8 sets bit 31, so the int32 container is negative — the - # unpack must still recover it (reads the 32-bit pattern as unsigned). + # Top nibble >= 8 makes the int32 container negative; unpack reads it unsigned. top = torch.tensor([[0, 0, 0, 0, 0, 0, 0, 8]], dtype=torch.int64) packed_top = pack_int32(top, num_bits=4) assert packed_top.item() == -(2**31) # 0x80000000 as signed int32 @@ -59,14 +33,7 @@ def test_pack_unpack_roundtrip(): assert packed.shape == (5, 4) assert torch.equal(unpack_int32(packed, num_bits=4), nibbles) - -# -------------------------------------------------------------------------- -# 2. Dequant math — symmetric offset-binary (nibble - bias) * scale. -# -------------------------------------------------------------------------- - def test_dequantize_symmetric_known_answer(): - # One row, one group of 8 (group_size=8). Unsigned nibbles minus bias 8 give - # the signed quantized values, times a per-group scale of 2.0. unsigned = torch.tensor([[8, 9, 7, 8, 10, 6, 8, 8]], dtype=torch.int64) signed = torch.tensor([[0, 1, -1, 0, 2, -2, 0, 0]], dtype=torch.float32) packed = pack_int32(unsigned, num_bits=4) @@ -79,8 +46,6 @@ def test_dequantize_symmetric_known_answer(): def test_dequantize_two_groups_broadcast(): - # in=16, group_size=8 -> two groups with distinct scales; check the scale - # broadcasts per-group along the input axis. unsigned = torch.full((1, 16), 8, dtype=torch.int64) unsigned[0, 0] = 9 # +1 in group 0 unsigned[0, 8] = 9 # +1 in group 1 @@ -94,11 +59,6 @@ def test_dequantize_two_groups_broadcast(): assert got[0, 8] == 5.0 assert got[0, 1] == 0.0 - -# -------------------------------------------------------------------------- -# 3. fake_quantize -> dequantize exactness (the golden's core invariant). -# -------------------------------------------------------------------------- - @pytest.mark.parametrize("group_size", [8, 16, -1]) def test_fake_quantize_dequantize_exact(group_size): torch.manual_seed(1) @@ -106,25 +66,17 @@ def test_fake_quantize_dequantize_exact(group_size): packed, scale, deq = fake_quantize_weight( w, num_bits=4, group_size=group_size, symmetric=True, ) - # Reconstructing from the on-disk tensors must reproduce the fake-quant bf16 - # result bit-for-bit (both compute q*scale in fp32 then cast to bf16). + # Both paths compute q*scale in fp32 then cast to bf16, so this is bit-exact. got = dequantize_weight( packed, scale, num_bits=4, group_size=group_size, symmetric=True, ) assert got.dtype == torch.bfloat16 assert torch.equal(got, deq) - # And the quantization is lossy but bounded (sanity: close to the original). assert (got.float() - w).abs().max() < 0.05 -# -------------------------------------------------------------------------- -# 4. The dequant-on-load streaming generator. -# -------------------------------------------------------------------------- - def _quant_components(base, w, cfg): - # Store the scale in bf16 (as a real compressed-tensors checkpoint does); the - # returned dequant is derived from that same bf16 scale, so it matches the - # stream's reconstruction bit-for-bit. + # Match checkpoint bf16 scales so stream reconstruction is bit-exact. packed, scale, deq = fake_quantize_weight( w, num_bits=cfg.num_bits, group_size=cfg.group_size, symmetric=cfg.symmetric, scale_dtype=torch.bfloat16, @@ -143,10 +95,9 @@ def test_stream_dequantizes_and_passes_through(): w_b = torch.randn(4, 16) * 0.1 comp_a, deq_a = _quant_components("layer.0.a_proj", w_a, cfg) comp_b, deq_b = _quant_components("layer.0.b_proj", w_b, cfg) - norm = torch.randn(8) # a non-quant key that must pass straight through + norm = torch.randn(8) - # Interleave/shuffle keys — the generator must reassemble by base name, - # independent of order. + # Reassembly is by base name, independent of stream order. stream = [ ("layer.0.a_proj.weight_scale", comp_a["layer.0.a_proj.weight_scale"]), ("layer.0.norm.weight", norm), @@ -157,7 +108,6 @@ def test_stream_dequantizes_and_passes_through(): ] out = dict(dequant_compressed_tensors_stream(iter(stream), cfg)) - # Exactly: two dequantized *.weight keys + the passthrough norm; no quant subkeys. assert set(out) == {"layer.0.a_proj.weight", "layer.0.b_proj.weight", "layer.0.norm.weight"} assert torch.equal(out["layer.0.a_proj.weight"], deq_a) assert torch.equal(out["layer.0.b_proj.weight"], deq_b) @@ -168,30 +118,21 @@ def test_stream_incomplete_group_raises(): cfg = CompressedTensorsQuantConfig(num_bits=4, group_size=16, symmetric=True) w = torch.randn(4, 16) * 0.1 comp, _ = _quant_components("x.proj", w, cfg) - # Only the packed tensor, no scale -> the group can never complete. stream = [("x.proj.weight_packed", comp["x.proj.weight_packed"])] with pytest.raises(ValueError, match="incomplete"): list(dequant_compressed_tensors_stream(iter(stream), cfg)) - -# -------------------------------------------------------------------------- -# 4b. dequant-on-load + packed-expert coexistence: keep_packed passes routed experts -# through raw while MLA/dense keys still dequantize (the streaming half of the -# mixed-load path; the GPU golden proves the packed params then load + run). -# -------------------------------------------------------------------------- - def test_stream_keep_packed_passthrough(): cfg = CompressedTensorsQuantConfig(num_bits=4, group_size=16, symmetric=True) - exp_base = "model.layers.1.mlp.experts.3.gate_proj" # a routed expert -> packed experts - mla_base = "model.layers.1.self_attn.o_proj" # an MLA weight -> dequant-on-load + exp_base = "model.layers.1.mlp.experts.3.gate_proj" + mla_base = "model.layers.1.self_attn.o_proj" comp_exp, _ = _quant_components(exp_base, torch.randn(8, 32) * 0.1, cfg) comp_mla, deq_mla = _quant_components(mla_base, torch.randn(4, 16) * 0.1, cfg) def keep_packed(base): return ".experts.3.gate_proj" in base - # Expert carries all three sub-keys (packed/scale/shape); MLA carries the two - # that complete a dequant (no shape) so no dangling buffer trips the end check. + # No dangling quant buffers should remain after mixed expert/MLA dequant. stream = [ (f"{exp_base}.weight_packed", comp_exp[f"{exp_base}.weight_packed"]), (f"{exp_base}.weight_scale", comp_exp[f"{exp_base}.weight_scale"]), @@ -201,28 +142,16 @@ def keep_packed(base): ] out = dict(dequant_compressed_tensors_stream(iter(stream), cfg, keep_packed=keep_packed)) - # Routed-expert sub-keys pass through RAW — packed int32 + scale + shape, and - # crucially NO collapsed ``.weight`` (they load into the packed params instead). assert out[f"{exp_base}.weight_packed"].dtype == torch.int32 assert f"{exp_base}.weight_scale" in out assert f"{exp_base}.weight_shape" in out assert f"{exp_base}.weight" not in out - # The MLA weight still collapses to one dequantized bf16 ``.weight`` (dequant-on-load). assert torch.equal(out[f"{mla_base}.weight"], deq_mla) assert f"{mla_base}.weight_packed" not in out - -# -------------------------------------------------------------------------- -# 4c. Pure-torch reference for the W4A16 kernel math (no GPU): the per-group -# ``(unpack - 8) * scale`` -> bf16 grouped GEMM the Triton kernel replicates, -# pinning the packed-K layout and the top-nibble (bit-31) sign case. -# -------------------------------------------------------------------------- - def test_grouped_gemm_reference_math_and_top_nibble(): torch.manual_seed(7) N, K, gs = 6, 32, 16 # two groups along the packed K axis - # Wide init so per-group amax uses the full nibble range -> some top nibbles - # land >= 8 (int32 container bit 31 set), exercising the sign path. w = torch.randn(N, K) * 0.5 packed, scale, deq = fake_quantize_weight( w, num_bits=4, group_size=gs, symmetric=True, scale_dtype=torch.bfloat16, @@ -230,25 +159,16 @@ def test_grouped_gemm_reference_math_and_top_nibble(): assert packed.shape == (N, K // 8) # packed along the last (input/K) axis assert (packed < 0).any(), "no negative container — top-nibble sign path untested" - # Reproduce the kernel's in-register arithmetic in pure torch: unpack the - # nibble, offset-binary subtract 8, scale per group (broadcast along K). nibbles = unpack_int32(packed, num_bits=4).to(torch.float32) # (N, K) unsigned scale_bc = scale.to(torch.float32).repeat_interleave(gs, dim=-1) # (N, K) manual_deq = ((nibbles - 8.0) * scale_bc).to(torch.bfloat16) assert torch.equal(manual_deq, deq) # kernel math == dequantize_weight - # Grouped GEMM equivalence: the packed path must produce the same y as feeding - # the bf16 dequant directly (both contract over the same bf16 weight values). x = torch.randn(4, K) y_manual = torch.einsum("tk,nk->tn", x, manual_deq.float()) y_deq = torch.einsum("tk,nk->tn", x, deq.float()) assert torch.equal(y_manual, y_deq) - -# -------------------------------------------------------------------------- -# 5. Config parsing. -# -------------------------------------------------------------------------- - def test_quant_config_from_hf_dict(): raw = { "format": "pack-quantized", From 8a7ce87fb07701c6c8b6724ae5d1345a53b6d09c Mon Sep 17 00:00:00 2001 From: Garv Ghai Date: Sat, 1 Aug 2026 18:41:17 +0000 Subject: [PATCH 7/9] Kimi-K2.7: register mla_absorb in the cache-manager backend test The weight-absorbed MLA work added a third ATTENTION_BACKENDS entry ("mla_absorb") but left test_registry_names asserting the exact pre-existing set, so the test has been failing since that change. Add the new backend to the expected set. --- test/modular/test_cache_manager_backends.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/modular/test_cache_manager_backends.py b/test/modular/test_cache_manager_backends.py index faa39f093..36039b418 100644 --- a/test/modular/test_cache_manager_backends.py +++ b/test/modular/test_cache_manager_backends.py @@ -77,4 +77,4 @@ def test_base_class_is_abstract(): def test_registry_names(): - assert set(ATTENTION_BACKENDS) == {"flashinfer", "dense_gen"} + assert set(ATTENTION_BACKENDS) == {"flashinfer", "dense_gen", "mla_absorb"} From 4a1e28340f2e1f6caf87d4a5b10b542fd9b644b1 Mon Sep 17 00:00:00 2001 From: Garv Ghai Date: Wed, 5 Aug 2026 01:21:05 +0000 Subject: [PATCH 8/9] Make the process-group timeout opt-in per deployment The 2h NCCL timeout applied to every multi-GPU model, so a genuinely hung collective sat 2h before aborting. Unset now means PyTorch's default; the Kimi TP8 configs opt in via dist_timeout_s, overridable per process with MSTAR_DIST_TIMEOUT_S. --- configs/kimi_k2_7_code_tp8.yaml | 3 +++ configs/kimi_k2_7_code_tp8_shm.yaml | 3 +++ docs/environment_variables.rst | 9 ++++++++ docs/serving.rst | 3 +++ mstar/conductor/conductor.py | 1 + mstar/distributed/communication.py | 32 +++++++++++++++++++++++++---- 6 files changed, 47 insertions(+), 4 deletions(-) diff --git a/configs/kimi_k2_7_code_tp8.yaml b/configs/kimi_k2_7_code_tp8.yaml index 3f9d48aeb..633a02081 100644 --- a/configs/kimi_k2_7_code_tp8.yaml +++ b/configs/kimi_k2_7_code_tp8.yaml @@ -1,6 +1,9 @@ model: "kimi_k2_7" # Fallback TP8 config for local Kimi-K2.7-Code snapshot; prefer the /dev/shm variant. max_seq_len: 8192 +# See kimi_k2_7_code_tp8_shm.yaml — a 1T INT4 load at TP8 outruns PyTorch's +# default process-group timeout; this path is slower still (disk, not tmpfs). +dist_timeout_s: 7200 model_kwargs: config_variant: k27_code checkpoint_path: /m-coriander/coriander/garv901/kimi_k2_7_code diff --git a/configs/kimi_k2_7_code_tp8_shm.yaml b/configs/kimi_k2_7_code_tp8_shm.yaml index bbd67205c..23a3396e7 100644 --- a/configs/kimi_k2_7_code_tp8_shm.yaml +++ b/configs/kimi_k2_7_code_tp8_shm.yaml @@ -1,6 +1,9 @@ model: "kimi_k2_7" # TP8 Kimi-K2.7-Code config using a RAM-backed /dev/shm checkpoint copy. max_seq_len: 8192 +# Loading ~600GB of INT4 experts across 8 ranks +# MSTAR_DIST_TIMEOUT_S overrides this. +dist_timeout_s: 7200 model_kwargs: config_variant: k27_code checkpoint_path: /dev/shm/kimi_k2_7_code diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index 72430b664..6eeed3d57 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -35,6 +35,15 @@ Communication - ``19000`` - Base of the deterministic entity-id → TCP port map (``api_server`` = base, ``conductor`` = base+1, ``worker_`` = base+100+rank). + * - ``MSTAR_DIST_TIMEOUT_S`` + - config's ``dist_timeout_s`` + - Timeout in seconds for the NCCL world group and its parallel + subgroups (:func:`mstar.distributed.communication.resolve_dist_timeout`). + Overrides the deployment config's ``dist_timeout_s``; with neither set, + PyTorch's default applies. Raise it only where weight load, JIT or + CUDA-graph capture can exceed that default (a 1T MoE at TP8 does) — + a hung collective takes correspondingly longer to abort. Must be set + before the conductor spawns workers, which inherit it. * - ``MSTAR_SHM_ARENA`` - ``0`` - SHM tensor-transport implementation. ``0``: per-uuid files. diff --git a/docs/serving.rst b/docs/serving.rst index 5b01c42df..080f49e16 100644 --- a/docs/serving.rst +++ b/docs/serving.rst @@ -178,6 +178,9 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key scoped to specific ``graph_walks`` and/or sharded with ``tp_size``. * - ``model_kwargs`` - *(optional)* Server-init model parameters (see below). + * - ``dist_timeout_s`` + - *(optional)* Timeout in seconds for the NCCL world group and its parallel + subgroups. Unset keeps PyTorch's default. Node names are model-specific — they are the keys of the model's ``get_node_engine_types`` (e.g. BAGEL's ``vit_encoder`` / ``vae_encoder`` / ``LLM``, diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index fe00e3c5d..5604f80f9 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -414,6 +414,7 @@ def _derive_worker_info(self): self.parallel_config = GlobalParallelConfig( worker_graphs=self.worker_graphs, worker_ids=self.worker_ids, + dist_timeout_s=self.model_config.get("dist_timeout_s"), ) def _launch_workers(self): diff --git a/mstar/distributed/communication.py b/mstar/distributed/communication.py index 3fd67061a..ae30b37a0 100644 --- a/mstar/distributed/communication.py +++ b/mstar/distributed/communication.py @@ -1,3 +1,4 @@ +import os from dataclasses import dataclass, field from datetime import timedelta from typing import Any @@ -5,6 +6,25 @@ import torch import torch.distributed as dist +DIST_TIMEOUT_ENV = "MSTAR_DIST_TIMEOUT_S" + + +def resolve_dist_timeout(dist_timeout_s: float | None = None) -> dict[str, timedelta]: + """Build the ``timeout`` kwarg for ``init_process_group`` / ``new_group``.""" + raw = os.environ.get(DIST_TIMEOUT_ENV, "").strip() + if raw: + try: + dist_timeout_s = float(raw) + except ValueError as exc: + raise ValueError( + f"{DIST_TIMEOUT_ENV} must be a number of seconds, got {raw!r}" + ) from exc + if dist_timeout_s is None: + return {} + if dist_timeout_s <= 0: + raise ValueError(f"Distributed timeout must be positive, got {dist_timeout_s}") + return {"timeout": timedelta(seconds=float(dist_timeout_s))} + class CommGroup: """A communication group over one axis of the worker device mesh. @@ -193,6 +213,8 @@ class WorkerParallelGroups: # projections). node_to_tp_group: dict[str, CommGroup] = field(default_factory=dict) node_to_sp_group: dict[str, CommGroup] = field(default_factory=dict) + # Process-group timeout in seconds, from the deployment config's + dist_timeout_s: float | None = None def add(self, node: str, comm_group: CommGroup): # disallow colocation of multiple comm groups on the same node @@ -234,13 +256,13 @@ def init_dist( if not self.any_parallelism: return - # Large-checkpoint load/JIT/capture can exceed the default NCCL timeout. + timeout_kwargs = resolve_dist_timeout(self.dist_timeout_s) dist.init_process_group( backend="nccl", init_method=init_method, world_size=self.num_workers, rank=self.global_rank, - timeout=timedelta(hours=2), + **timeout_kwargs, ) # One subgroup per distinct rank tuple across BOTH mesh axes — @@ -251,7 +273,7 @@ def init_dist( rank_tuple_to_pg: dict[tuple[int, ...], "dist.ProcessGroup"] = {} for rank_tuple in self.world_parallel_groups: rank_tuple_to_pg[rank_tuple] = dist.new_group( - ranks=list(rank_tuple), timeout=timedelta(hours=2) + ranks=list(rank_tuple), **timeout_kwargs ) seen: set[int] = set() @@ -314,7 +336,8 @@ class GlobalParallelConfig: def __init__( # leaving type annotation as Any due to circular import self, worker_graphs: dict[str, Any], - worker_ids: list[str] + worker_ids: list[str], + dist_timeout_s: float | None = None, ): self.num_workers = len(worker_ids) any_parallelism = any( @@ -342,6 +365,7 @@ def __init__( global_rank=i, num_workers=self.num_workers, any_parallelism=any_parallelism, world_parallel_groups=world_parallel_groups, + dist_timeout_s=dist_timeout_s, ) for i, wid in enumerate(worker_ids) } From 58ede843da681f0724780bed8ec9e47c4c69f22e Mon Sep 17 00:00:00 2001 From: Garv Ghai Date: Sat, 8 Aug 2026 16:56:18 +0000 Subject: [PATCH 9/9] Addressing comments: QuantizationData abstraction, MLA capture guard, --model-path --- configs/kimi_k2_7_code_tp8.yaml | 5 +- configs/kimi_k2_7_code_tp8_shm.yaml | 5 +- configs/kimi_k2_7_repro.yaml | 15 -- configs/kimi_k2_7_tp2.yaml | 15 -- configs/synthetic/kimi_k2_7_repro.yaml | 17 +++ configs/synthetic/kimi_k2_7_tp2.yaml | 17 +++ docs/serving.rst | 19 +++ mstar/api_server/entrypoint.py | 29 +++- mstar/cli/main.py | 8 ++ mstar/engine/cache_manager.py | 100 ++++++++++--- mstar/engine/cuda_graph_runner.py | 43 ++++-- .../model/components/quantization/__init__.py | 16 +++ .../quantization/compressed_tensors.py} | 66 ++++++++- .../components/quantization/marlin_moe.py | 33 ++++- mstar/model/kimi_k2_7/_testing.py | 9 +- mstar/model/kimi_k2_7/components/attention.py | 13 ++ mstar/model/kimi_k2_7/components/moe.py | 35 +++-- mstar/model/kimi_k2_7/config.py | 2 +- mstar/model/kimi_k2_7/kimi_model.py | 8 +- mstar/model/kimi_k2_7/weight_loader.py | 4 +- mstar/utils/fused_moe/runner.py | 117 ++++++++------- mstar/utils/quantization.py | 90 ++++++++++++ .../integration/test_kimi_mla_absorb_paged.py | 24 ++++ .../test_kimi_moe_inkernel_dequant.py | 6 +- test/integration/test_kimi_serve_e2e.py | 8 +- test/integration/test_marlin_kernels.py | 3 +- test/modular/test_cache_manager_backends.py | 85 ++++++++++- ...mi_quant.py => test_compressed_tensors.py} | 4 +- test/modular/test_quantization_types.py | 136 ++++++++++++++++++ 29 files changed, 773 insertions(+), 159 deletions(-) delete mode 100644 configs/kimi_k2_7_repro.yaml delete mode 100644 configs/kimi_k2_7_tp2.yaml create mode 100644 configs/synthetic/kimi_k2_7_repro.yaml create mode 100644 configs/synthetic/kimi_k2_7_tp2.yaml rename mstar/model/{kimi_k2_7/quantization.py => components/quantization/compressed_tensors.py} (65%) create mode 100644 mstar/utils/quantization.py rename test/modular/{test_kimi_quant.py => test_compressed_tensors.py} (99%) create mode 100644 test/modular/test_quantization_types.py diff --git a/configs/kimi_k2_7_code_tp8.yaml b/configs/kimi_k2_7_code_tp8.yaml index 633a02081..e05bd5f34 100644 --- a/configs/kimi_k2_7_code_tp8.yaml +++ b/configs/kimi_k2_7_code_tp8.yaml @@ -4,9 +4,12 @@ max_seq_len: 8192 # See kimi_k2_7_code_tp8_shm.yaml — a 1T INT4 load at TP8 outruns PyTorch's # default process-group timeout; this path is slower still (disk, not tmpfs). dist_timeout_s: 7200 +# Point at your checkpoint at launch — this config hardcodes no path: +# mstar-serve --config configs/kimi_k2_7_code_tp8.yaml \ +# --model-path /path/to/Kimi-K2.7-Code +# Omit --model-path to pull moonshotai/Kimi-K2.7-Code from HuggingFace. model_kwargs: config_variant: k27_code - checkpoint_path: /m-coriander/coriander/garv901/kimi_k2_7_code tokenizer_mode: hf kv_cache: max_num_pages: 512 diff --git a/configs/kimi_k2_7_code_tp8_shm.yaml b/configs/kimi_k2_7_code_tp8_shm.yaml index 23a3396e7..3298aa027 100644 --- a/configs/kimi_k2_7_code_tp8_shm.yaml +++ b/configs/kimi_k2_7_code_tp8_shm.yaml @@ -4,9 +4,12 @@ max_seq_len: 8192 # Loading ~600GB of INT4 experts across 8 ranks # MSTAR_DIST_TIMEOUT_S overrides this. dist_timeout_s: 7200 +# Copy the checkpoint into /dev/shm first, then point at it at launch: +# mstar serve kimi_k2_7 --model-path /dev/shm/kimi_k2_7_code +# mstar-serve --config configs/kimi_k2_7_code_tp8_shm.yaml \ +# --model-path /dev/shm/kimi_k2_7_code model_kwargs: config_variant: k27_code - checkpoint_path: /dev/shm/kimi_k2_7_code tokenizer_mode: hf kv_cache: max_num_pages: 512 diff --git a/configs/kimi_k2_7_repro.yaml b/configs/kimi_k2_7_repro.yaml deleted file mode 100644 index 358bcd43d..000000000 --- a/configs/kimi_k2_7_repro.yaml +++ /dev/null @@ -1,15 +0,0 @@ -model: "kimi_k2_7" -# Reduced synthetic serve config. Generate tools/kimi_goldens/repro/checkpoint first. -max_seq_len: 512 -model_kwargs: - config_variant: reduced - checkpoint_path: tools/kimi_goldens/repro/checkpoint - tokenizer_mode: byte -kv_cache: - max_num_pages: 256 - page_size: 128 -node_groups: - - node_names: [LLM] - ranks: [0] - tp_size: 1 - graph_walks: [prefill, decode] diff --git a/configs/kimi_k2_7_tp2.yaml b/configs/kimi_k2_7_tp2.yaml deleted file mode 100644 index 01ae87549..000000000 --- a/configs/kimi_k2_7_tp2.yaml +++ /dev/null @@ -1,15 +0,0 @@ -model: "kimi_k2_7" -# Reduced synthetic TP=2 config. Generate tools/kimi_goldens/repro/checkpoint first. -max_seq_len: 512 -model_kwargs: - config_variant: reduced - checkpoint_path: tools/kimi_goldens/repro/checkpoint - tokenizer_mode: byte -kv_cache: - max_num_pages: 256 - page_size: 128 -node_groups: - - node_names: [LLM] - ranks: [0, 1] - tp_size: 2 - graph_walks: [prefill, decode] diff --git a/configs/synthetic/kimi_k2_7_repro.yaml b/configs/synthetic/kimi_k2_7_repro.yaml new file mode 100644 index 000000000..bc3345d35 --- /dev/null +++ b/configs/synthetic/kimi_k2_7_repro.yaml @@ -0,0 +1,17 @@ +model: "kimi_k2_7" +# Reduced synthetic serve config — random weights, NOT a production deployment. +# Generate the checkpoint first, then point at it: +# mstar-serve --config configs/synthetic/kimi_k2_7_repro.yaml \ +# --model-path tools/kimi_goldens/repro/checkpoint +max_seq_len: 512 +model_kwargs: + config_variant: reduced + tokenizer_mode: byte +kv_cache: + max_num_pages: 256 + page_size: 128 +node_groups: + - node_names: [LLM] + ranks: [0] + tp_size: 1 + graph_walks: [prefill, decode] diff --git a/configs/synthetic/kimi_k2_7_tp2.yaml b/configs/synthetic/kimi_k2_7_tp2.yaml new file mode 100644 index 000000000..17bd7c82f --- /dev/null +++ b/configs/synthetic/kimi_k2_7_tp2.yaml @@ -0,0 +1,17 @@ +model: "kimi_k2_7" +# Reduced synthetic TP=2 config — random weights, NOT a production deployment. +# Generate the checkpoint first, then point at it: +# mstar-serve --config configs/synthetic/kimi_k2_7_tp2.yaml \ +# --model-path tools/kimi_goldens/repro/checkpoint +max_seq_len: 512 +model_kwargs: + config_variant: reduced + tokenizer_mode: byte +kv_cache: + max_num_pages: 256 + page_size: 128 +node_groups: + - node_names: [LLM] + ranks: [0, 1] + tp_size: 2 + graph_walks: [prefill, decode] diff --git a/docs/serving.rst b/docs/serving.rst index 080f49e16..b63d9af6a 100644 --- a/docs/serving.rst +++ b/docs/serving.rst @@ -44,6 +44,10 @@ mstar serve * - ``--cache-dir`` - HF default - HuggingFace weight cache directory. + * - ``--model-path`` + - registry default + - Local checkpoint directory or HF repo id to load weights from. Use this + instead of hardcoding a path in the config YAML. * - ``--tensor-comm-protocol`` - ``SHM`` - Tensor transport: ``SHM`` (safe single-node default), ``TCP``, or ``RDMA``. @@ -122,6 +126,16 @@ mstar-serve * - ``--cache-dir`` - HF default - HuggingFace weight cache directory. + * - ``--model-path`` + - registry default + - Local checkpoint directory or HF repo id to load weights from. Applies to + any model and overrides the registry default, so a config YAML never has + to hardcode one machine's filesystem layout: + + .. code-block:: bash + + mstar-serve --config configs/kimi_k2_7_code_tp8_shm.yaml \ + --model-path /dev/shm/kimi_k2_7_code * - ``--socket-path-prefix`` - ``/tmp/mstar`` - ZMQ IPC socket prefix (shared with conductor/workers). @@ -211,6 +225,11 @@ Because placement is config-only, the *same* model code runs single-GPU or fully disaggregated. ``configs/`` ships several layouts per model (``*_single_gpu``, ``*_colocated``, ``*_pd_disaggregated``, ``*_cfg_parallel``, …). +``configs/synthetic/`` holds reduced, randomly-initialized deployments used to +exercise the serving path without real weights (shape/plumbing checks, TP +sanity). They are **not** production configs — they load a generated checkpoint +and emit meaningless tokens. + **Tensor parallelism.** Shard a node across GPUs with ``tp_size`` and that many ``ranks``: .. code-block:: yaml diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 3fc97ad4d..160aa7f1a 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -52,6 +52,19 @@ def _detect_modality(filename: str) -> str: # Conductor process target (top-level for picklability with spawn) # ------------------------------------------------------------------ +def _resolve_model_path(model_name: str, model_path: str | None) -> str: + """Where to load weights from: ``--model-path`` if given, else the registry default. + + ``--model-path`` takes a local directory or an HF repo id and applies to any + model, so a deployment config never has to hardcode one machine's filesystem + layout. Both the API-server-side model instance and the conductor-side one + resolve through here, so they cannot diverge. + """ + if model_path: + return model_path + return HF_MODELS.get(model_name, {}).get("model_path_hf", "") + + def _conductor_process_target( model_name: str, config_path: str, @@ -61,7 +74,8 @@ def _conductor_process_target( log_level: str = "INFO", cache_dir: str | None = None, tensor_comm_protocol=CommProtocol.RDMA, - tcp_transfer_device="" + tcp_transfer_device="", + model_path: str | None = None, ): """Runs DummyConductor.run() in a spawned process.""" logging.basicConfig( @@ -93,7 +107,7 @@ def _conductor_process_target( ) model = get_model_class(model_name)( - model_path_hf=HF_MODELS.get(model_name, {}).get("model_path_hf", ""), + model_path_hf=_resolve_model_path(model_name, model_path), cache_dir=cache_dir, **yaml_model_kwargs, ) @@ -866,6 +880,12 @@ def main(argv: list[str] | None = None): "--cache-dir", type=str, default=None, help="Directory for caching downloaded HuggingFace model files", ) + parser.add_argument( + "--model-path", type=str, default=None, + help="Where to load weights from — a local checkpoint directory or an HF " + "repo id. Overrides the model's registry default so deployment " + "configs need not hardcode a filesystem path.", + ) parser.add_argument( "--log-level", type=str, default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"], @@ -900,7 +920,7 @@ def main(argv: list[str] | None = None): # (tokenization only — no GPU weights needed) from mstar.model.registry import get_model_class model = get_model_class(model_name)( - model_path_hf=HF_MODELS.get(model_name, {}).get("model_path_hf", ""), + model_path_hf=_resolve_model_path(model_name, args.model_path), cache_dir=args.cache_dir, **yaml_model_kwargs, ) @@ -932,7 +952,8 @@ def main(argv: list[str] | None = None): args.log_level, args.cache_dir, CommProtocol(args.tensor_comm_protocol), - args.tcp_transfer_device + args.tcp_transfer_device, + args.model_path, ), ) conductor_proc.start() diff --git a/mstar/cli/main.py b/mstar/cli/main.py index 22a59acd9..45d734663 100644 --- a/mstar/cli/main.py +++ b/mstar/cli/main.py @@ -38,6 +38,7 @@ "whisper_large": "whisper_large.yaml", "higgs_audio": "higgs_audio.yaml", "wan22": "wan22.yaml", + "kimi_k2_7": "kimi_k2_7_code_tp8_shm.yaml", } @@ -149,6 +150,8 @@ def _serve(args: argparse.Namespace) -> None: ] if args.cache_dir: argv += ["--cache-dir", args.cache_dir] + if args.model_path: + argv += ["--model-path", args.model_path] if args.log_stats: argv += ["--log-stats"] if args.log_stats_file: @@ -172,6 +175,11 @@ def build_parser() -> argparse.ArgumentParser: serve.add_argument("--gpus", default=None, help="CUDA_VISIBLE_DEVICES, e.g. '0' or '0,1,2'") serve.add_argument("--config", default=None, help="override the default config (path to YAML)") serve.add_argument("--cache-dir", default=None, help="HuggingFace weight cache directory") + serve.add_argument( + "--model-path", default=None, + help="local checkpoint directory or HF repo id to load weights from " + "(overrides the model's registry default)", + ) serve.add_argument("--socket-path-prefix", default=None, help="ZMQ IPC socket prefix") serve.add_argument("--upload-dir", default=None, help="temp dir for uploaded media") serve.add_argument( diff --git a/mstar/engine/cache_manager.py b/mstar/engine/cache_manager.py index 209fb1525..5d586cc73 100644 --- a/mstar/engine/cache_manager.py +++ b/mstar/engine/cache_manager.py @@ -78,6 +78,36 @@ class PlanCacheKey(NamedTuple): dtype: torch.dtype +@dataclass +class MlaRequestSlice: + """One request's slice of an absorbed-MLA SDPA plan. + + ``q_start``/``seq_len`` locate this request's rows in the packed query batch; + ``total_len`` is its context length after this step (so the causal mask knows + how many cached tokens precede the new ones); ``page_indices`` gathers its + latent pages. + """ + q_start: int + seq_len: int + total_len: int + page_indices: torch.Tensor + + +@dataclass +class MlaSdpaPlan: + """Absorbed-MLA fallback plan: latent scatter indices + per-request gather layout. + + Built only when the FlashInfer MLA kernel cannot serve the configured latent + dims (see :func:`mla_kernel_available_for`); the kernel path plans a + :class:`FlashInferMLAWrapper` instead. For an ``mla_absorb`` label exactly one + of ``_PlanState.wrapper`` / ``_PlanState.mla`` is set — ``run_attention_mla`` + uses that to choose a path, so leaving a stale value in either is a bug. + """ + token_to_page: torch.Tensor + token_to_cache: torch.Tensor + requests: list[MlaRequestSlice] + + @dataclass class _PlanState: """Pre-computed state from plan_attention/plan_rope for a single cache label. @@ -127,8 +157,9 @@ class _PlanState: # segment over its contiguous frozen prefix. None on paged plans, which # keep the FlashInfer path. See DenseGenCacheManager._build_dense_gen_plan. dense_gen: dict | None = None - # MLA absorb fallback plan: latent scatter indices and per-request gather layout. - mla: dict | None = None + # MLA absorb fallback plan: latent scatter indices and per-request gather + # layout. Mutually exclusive with ``wrapper`` (see MlaSdpaPlan). + mla: "MlaSdpaPlan | None" = None class WorkspaceBufferManager: @@ -1519,6 +1550,27 @@ def _mla_kernel_available(ckv: int, kpe: int, sm_major: int) -> bool: return True +def mla_kernel_available_for(kv_cache_config: KVCacheConfig, device) -> bool: + """Whether ``kv_cache_config``'s latent dims can use the FlashInfer MLA kernel. + + The single predicate behind both the runtime plan (:class:`MlaAbsorbCacheManager`) + and the capture decision (``cuda_graph_runner.mla_absorb_capture_blocked``). They + must not drift: a capture that assumes the kernel while the plan takes SDPA builds + a paged wrapper the 4-D latent cache cannot serve. + + Total by construction — a non-CUDA device returns False rather than raising, so the + absorbed-SDPA fallback is reachable on CPU. ``_mla_kernel_available`` explains why + this must be decided *before* any kernel is constructed. + """ + ckv = kv_cache_config.mla_ckv_dim + if ckv is None: + return False + if not torch.cuda.is_available() or torch.device(device).type != "cuda": + return False + sm_major = torch.cuda.get_device_capability(device)[0] + return _mla_kernel_available(ckv, kv_cache_config.head_dim - ckv, sm_major) + + class MlaAbsorbCacheManager(FlashInferCacheManager): """Paged-cache backend for weight-absorbed MLA. @@ -1549,8 +1601,7 @@ def plan_attention( page_size = cfg.page_size ckv = cfg.mla_ckv_dim kpe = (cfg.head_dim - ckv) if ckv is not None else None - sm_major = torch.cuda.get_device_capability(self.device)[0] - use_kernel = ckv is not None and _mla_kernel_available(ckv, kpe, sm_major) + use_kernel = mla_kernel_available_for(cfg, self.device) if use_kernel: qo_indptr_list = [0] @@ -1599,7 +1650,7 @@ def plan_attention( else: token_to_page: list[int] = [] token_to_cache: list[int] = [] - requests: list[dict] = [] + requests: list[MlaRequestSlice] = [] q_start = 0 for i, rid in enumerate(self.request_ids): state = self._get_state(rid, effective_label) @@ -1618,25 +1669,30 @@ def plan_attention( token_to_page.append(page_indices[g // page_size]) token_to_cache.append(g % page_size) - requests.append({ - "q_start": q_start, - "seq_len": sl, - "total_len": total_len, - "page_indices": torch.tensor( + requests.append(MlaRequestSlice( + q_start=q_start, + seq_len=sl, + total_len=total_len, + page_indices=torch.tensor( page_indices, dtype=torch.long, device=self.device ), - }) + )) q_start += sl - ps.mla = { - "token_to_page": torch.tensor( + ps.mla = MlaSdpaPlan( + token_to_page=torch.tensor( token_to_page, dtype=torch.long, device=self.device ), - "token_to_cache": torch.tensor( + token_to_cache=torch.tensor( token_to_cache, dtype=torch.long, device=self.device ), - "requests": requests, - } + requests=requests, + ) + # Clear any wrapper a CUDA-graph capture injected: run_attention_mla + # picks its path with ``ps.wrapper is not None``, and a paged wrapper + # cannot read the 4-D latent cache. Mirrors how the paged path clears + # ``dense_gen`` in DenseGenCacheManager. + ps.wrapper = None # Keep advance_seq_lens / flush_to_store aligned with the base manager. ps.seq_lens = seq_lens @@ -1683,7 +1739,7 @@ def run_attention_mla( mla = ps.mla assert mla is not None - latent_cache[mla["token_to_page"], mla["token_to_cache"]] = latent.to( + latent_cache[mla.token_to_page, mla.token_to_cache] = latent.to( latent_cache.dtype ) @@ -1692,13 +1748,13 @@ def run_attention_mla( query_all = torch.cat([q_nope, q_pe], dim=-1) out = torch.empty(T, H, L, dtype=q_nope.dtype, device=q_nope.device) - for req in mla["requests"]: - q_start = req["q_start"] - sl = req["seq_len"] - total_len = req["total_len"] + for req in mla.requests: + q_start = req.q_start + sl = req.seq_len + total_len = req.total_len # Mirror the dense-gen page gather for the fallback. - gathered = latent_cache[req["page_indices"]].reshape( + gathered = latent_cache[req.page_indices].reshape( -1, latent_cache.shape[-1] )[:total_len] key = gathered diff --git a/mstar/engine/cuda_graph_runner.py b/mstar/engine/cuda_graph_runner.py index 7637ad01a..ed16b3579 100644 --- a/mstar/engine/cuda_graph_runner.py +++ b/mstar/engine/cuda_graph_runner.py @@ -30,6 +30,7 @@ BatchedCacheManager, WorkspaceBufferManager, create_cache_manager, + mla_kernel_available_for, ) from mstar.engine.cuda_graph_config import ( BasicBatchedCudaGraphConfig, @@ -56,6 +57,21 @@ DEFAULT_AR_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32, 64] +def mla_absorb_capture_blocked( + kv_cache_config: KVCacheConfig, device +) -> str | None: + if kv_cache_config.attention_backend != "mla_absorb": + return None + if mla_kernel_available_for(kv_cache_config, device): + return None + ckv = kv_cache_config.mla_ckv_dim + return ( + "attention_backend='mla_absorb' without the FlashInfer MLA kernel " + f"(mla_ckv_dim={ckv}, kpe={None if ckv is None else kv_cache_config.head_dim - ckv}, " + f"device={device}); the absorbed-SDPA fallback runs eager and plans no wrapper" + ) + + @dataclass class DummyCaptureInput: tensors: dict[str, list[torch.Tensor]] # {tensor_name: [tensor(s)]} @@ -258,6 +274,14 @@ def warmup_and_capture(self) -> None: self.submodule_name) return + blocked = mla_absorb_capture_blocked(self.kv_cache_config, self.device) + if blocked is not None: + logger.info( + "Skipping CUDA graph capture for %s: %s. This node serves eagerly.", + self.submodule_name, blocked, + ) + return + if not hasattr(self.submodule, 'forward_batched'): logger.info("Submodule %s does not support batched forward, " "skipping CUDA graph capture", self.submodule_name) @@ -326,7 +350,7 @@ def _create_persistent_wrappers( can run on plan_stream concurrently with replay(slot 0) on default_stream without racing on the wrapper's persistent state. """ - from mstar.engine.cache_manager import _mla_kernel_available, _PlanState + from mstar.engine.cache_manager import _PlanState from mstar.utils.flashinfer_utils import ( FlashInferDecodeWrapper, FlashInferMLAWrapper, @@ -338,15 +362,16 @@ def _create_persistent_wrappers( cfg = self.kv_cache_config # Only the FlashInfer MLA path is capturable; absorbed SDPA stays eager. - use_mla_kernel = ( - cfg.attention_backend == "mla_absorb" - and cfg.mla_ckv_dim is not None - and _mla_kernel_available( - cfg.mla_ckv_dim, - cfg.head_dim - cfg.mla_ckv_dim, - torch.cuda.get_device_capability(self.device)[0], + # warmup_and_capture already cancelled capture in that case, so reaching + # here with mla_absorb-and-no-kernel means the two decisions drifted — + # refuse rather than fall through and build a paged wrapper that cannot + # read the 4-D latent cache. + blocked = mla_absorb_capture_blocked(cfg, self.device) + if blocked is not None: + raise RuntimeError( + f"_create_persistent_wrappers called for an uncapturable config: {blocked}" ) - ) + use_mla_kernel = cfg.attention_backend == "mla_absorb" # Allocate workspace buffer for CUDA graph wrappers. # Each (label, slot) gets its own workspace — slots must NOT share diff --git a/mstar/model/components/quantization/__init__.py b/mstar/model/components/quantization/__init__.py index 93d1b85dd..2a39f24db 100644 --- a/mstar/model/components/quantization/__init__.py +++ b/mstar/model/components/quantization/__init__.py @@ -3,10 +3,26 @@ FusedMoEQuantizeMethod, process_weights_after_loading, ) +from mstar.model.components.quantization.compressed_tensors import ( + CompressedTensorsQuantConfig, + dequant_compressed_tensors_stream, + dequantize_weight, + pack_int32, + unpack_int32, +) from mstar.model.components.quantization.marlin_moe import MarlinMoEMethod +from mstar.utils.quantization import QuantizationData, QuantizationType, W4A16Data __all__ = [ + "CompressedTensorsQuantConfig", "FusedMoEQuantizeMethod", "MarlinMoEMethod", + "QuantizationData", + "QuantizationType", + "W4A16Data", + "dequant_compressed_tensors_stream", + "dequantize_weight", + "pack_int32", "process_weights_after_loading", + "unpack_int32", ] diff --git a/mstar/model/kimi_k2_7/quantization.py b/mstar/model/components/quantization/compressed_tensors.py similarity index 65% rename from mstar/model/kimi_k2_7/quantization.py rename to mstar/model/components/quantization/compressed_tensors.py index 7f2ffef22..ad96c5ca9 100644 --- a/mstar/model/kimi_k2_7/quantization.py +++ b/mstar/model/components/quantization/compressed_tensors.py @@ -1,7 +1,17 @@ -"""Compressed-tensors INT4 W4A16 helpers for Kimi-K2.7 weights. +"""Reader for the compressed-tensors checkpoint format (neuralmagic / vLLM). + +Model-agnostic: this is a *serialization format*, not one model's quirk. Kimi-K2.7 +is simply the first mstar model whose checkpoint ships in it. Packed values are stored low-order-first in int32 containers. Symmetric INT4 is offset-binary, so dequant subtracts 8 to match vLLM's ``uint4b8`` layout. + +:class:`CompressedTensorsQuantConfig` describes how a checkpoint was *written*; +:meth:`CompressedTensorsQuantConfig.moe_quant_data` translates that into the +:class:`~mstar.utils.quantization.QuantizationData` a specific kernel call needs. +:attr:`~CompressedTensorsQuantConfig.quant_type` is the only place ``num_bits`` is +mapped to a supported scheme; :meth:`~CompressedTensorsQuantConfig.ensure_kernel_support` +is how callers reject a width no kernel implements. """ from __future__ import annotations @@ -10,6 +20,8 @@ import torch +from mstar.utils.quantization import QuantizationData, QuantizationType, W4A16Data + _PACKED = ".weight_packed" _SCALE = ".weight_scale" _ZERO_POINT = ".weight_zero_point" @@ -33,6 +45,58 @@ class CompressedTensorsQuantConfig: def pack_factor(self) -> int: return 32 // self.num_bits + @property + def quant_type(self) -> QuantizationType | None: + """The scheme mstar's kernels implement for this checkpoint, else ``None``. + + The single ``num_bits`` -> scheme decision. ``None`` means "no kernel for + this width"; call :meth:`ensure_kernel_support` to turn that into an error. + """ + return QuantizationType.W4A16 if self.num_bits == 4 else None + + def ensure_kernel_support(self) -> QuantizationType: + """Return this checkpoint's scheme, raising if mstar has no kernel for it. + + Call before allocating packed parameters or selecting a backend. The + packed-expert kernels hardcode 4-bit nibble extraction, so an INT8 + checkpoint would otherwise allocate self-consistent shapes, load without + complaint, and return wrong numbers — this turns that into a load-time + error naming the width the checkpoint declared. + """ + quant_type = self.quant_type + if quant_type is None: + raise ValueError( + f"compressed-tensors checkpoint declares num_bits={self.num_bits}, which " + "mstar does not implement (supported: 4-bit / W4A16). The fused-MoE and " + "Marlin kernels are INT4-only." + ) + return quant_type + + def moe_quant_data( + self, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w1_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, + ) -> QuantizationData: + """Companion data for one stacked routed-expert GEMM under this config. + + Call per dispatch rather than caching on the module: the scales are + ``nn.Parameter`` s that ``Module._apply`` rebuilds on ``.to(device)``, so + binding them at call time keeps the returned object from ever holding a + stale tensor. + """ + quant_type = self.ensure_kernel_support() + if quant_type is QuantizationType.W4A16: + return W4A16Data( + w1_scale=w1_scale, + w2_scale=w2_scale, + group_size=self.group_size, + w1_zp=None if self.symmetric else w1_zp, + w2_zp=None if self.symmetric else w2_zp, + ) + raise ValueError(f"No routed-expert quantization data for {quant_type}") + @classmethod def from_hf_config_dict( cls, quant: dict | None diff --git a/mstar/model/components/quantization/marlin_moe.py b/mstar/model/components/quantization/marlin_moe.py index e7000a296..fd707d75c 100644 --- a/mstar/model/components/quantization/marlin_moe.py +++ b/mstar/model/components/quantization/marlin_moe.py @@ -1,18 +1,26 @@ """Marlin W4A16 backend for routed-expert MoE GEMMs.""" from __future__ import annotations +from typing import TYPE_CHECKING + import torch from mstar.utils.marlin import ops as marlin_ops +from mstar.utils.quantization import QuantizationType + +if TYPE_CHECKING: + from mstar.model.components.quantization.compressed_tensors import ( + CompressedTensorsQuantConfig, + ) class MarlinMoEMethod: """Marlin routed-expert GEMM backend (symmetric INT4, group-wise). - Stateful: :meth:`prepare` repacks the loaded packed experts into Marlin - layout and stores them (plus the workspace) on the instance; the source - packed params can then be freed by the owning block. :meth:`apply` runs the - two Marlin GEMMs. + A kernel *backend*, not a quantization descriptor: stateful, and it owns + Marlin-layout weights after :meth:`prepare` repacks the loaded packed experts + (the source packed params can then be freed by the owning block). + :meth:`apply` runs the two Marlin GEMMs. """ def __init__(self, *, num_bits: int = 4, group_size: int = 32) -> None: @@ -28,6 +36,23 @@ def __init__(self, *, num_bits: int = 4, group_size: int = 32) -> None: self.w2_scale: torch.Tensor | None = None self.workspace: torch.Tensor | None = None + @classmethod + def from_quant_config( + cls, quant_config: "CompressedTensorsQuantConfig" + ) -> "MarlinMoEMethod": + """Build from the checkpoint descriptor, so the bit width is read forward + from ``num_bits`` rather than reconstructed backwards from a pack factor.""" + quant_type = quant_config.ensure_kernel_support() + if quant_type is not QuantizationType.W4A16: + raise ValueError( + f"MarlinMoEMethod supports {QuantizationType.W4A16} only, got {quant_type}" + ) + return cls(num_bits=quant_config.num_bits, group_size=quant_config.group_size) + + @property + def quant_type(self) -> QuantizationType: + return QuantizationType.W4A16 + def prepare( self, w13_packed: torch.Tensor, diff --git a/mstar/model/kimi_k2_7/_testing.py b/mstar/model/kimi_k2_7/_testing.py index b814f089c..bcc741a0d 100644 --- a/mstar/model/kimi_k2_7/_testing.py +++ b/mstar/model/kimi_k2_7/_testing.py @@ -2,20 +2,21 @@ NOT part of the serving path. These helpers exist only to let the test suite fabricate a synthetic quantized checkpoint and its exact bf16 reference, so a -golden can assert the real load path (``quantization.py`` + -``weight_loader.py``) reproduces the reference bit-for-bit. Nothing here is +golden can assert the real load path (:mod:`mstar.model.components.quantization` ++ ``weight_loader.py``) reproduces the reference bit-for-bit. Nothing here is imported by the model or the loader at serve time. The real load-path primitives (``pack_int32`` / ``unpack_int32`` / ``dequantize_weight`` / ``dequant_compressed_tensors_stream`` / -``CompressedTensorsQuantConfig``) live in ``quantization.py``; this module builds +``CompressedTensorsQuantConfig``) live in +:mod:`mstar.model.components.quantization.compressed_tensors`; this module builds on them. """ from __future__ import annotations import torch -from mstar.model.kimi_k2_7.quantization import dequantize_weight, pack_int32 +from mstar.model.components.quantization import dequantize_weight, pack_int32 def fake_quantize_weight( diff --git a/mstar/model/kimi_k2_7/components/attention.py b/mstar/model/kimi_k2_7/components/attention.py index 27100b5b5..e4c44c8b1 100644 --- a/mstar/model/kimi_k2_7/components/attention.py +++ b/mstar/model/kimi_k2_7/components/attention.py @@ -86,6 +86,19 @@ def forward( ) -> torch.Tensor: if self.mla_absorb: return self._forward_absorbed(hidden_states, cache_handle, position_ids) + + # Naive/materialized MLA — the reduced-capability fallback (see the + # ``mla_absorb`` note in config.py), not a co-equal path. It projects the + # latent up to full per-head K/V and uses the standard paged MHA + # interface, so it works on any dims/hardware, at the cost of caching + # H*(Dnope+Dv) per token instead of one shared latent. Kept for parity + # tests and for configs the absorbed path's kernels cannot serve. + # + # Two numerical quirks follow from reusing that MHA interface: q/k/v are + # padded to ``padded_head_dim`` (FlashInfer paged kernels only accept + # 64/128/256), and because ``run_attention`` then scales by + # 1/sqrt(padded_head_dim), DeepSeek's intended qk_head_dim**-0.5 * mscale**2 + # is folded into q via ``softmax_scale_boost`` below. num_tokens = hidden_states.shape[0] h = self.num_heads diff --git a/mstar/model/kimi_k2_7/components/moe.py b/mstar/model/kimi_k2_7/components/moe.py index f744f681d..f7a195735 100644 --- a/mstar/model/kimi_k2_7/components/moe.py +++ b/mstar/model/kimi_k2_7/components/moe.py @@ -164,6 +164,10 @@ def __init__( self.packed_experts = ( config.quantization_config is not None and config.moe_in_kernel_dequant ) + # Hold the checkpoint descriptor rather than copying its fields out: it is + # the single source for group_size / pack_factor / symmetric, and it builds + # the per-dispatch QuantizationData. + self.quant_config = config.quantization_config if self.packed_experts else None self.quant_kernel = getattr(config, "quant_kernel", "auto") self._marlin_method = None self._use_marlin = False @@ -182,11 +186,9 @@ def __init__( self.experts = nn.Module() if self.packed_experts: - qc = config.quantization_config - self.group_size = qc.group_size - self.pack_factor = qc.pack_factor # 8 for INT4 - self.symmetric = qc.symmetric - hidden, gs, pf = config.hidden_size, self.group_size, self.pack_factor + qc = self.quant_config + qc.ensure_kernel_support() + hidden, gs, pf = config.hidden_size, qc.group_size, qc.pack_factor assert hidden % pf == 0 and hidden % gs == 0, ( f"hidden {hidden} must divide pack_factor {pf} and group_size {gs}" ) @@ -241,7 +243,7 @@ def _attach_expert_weight_loaders(self) -> None: full_inter = self.moe_intermediate_size if self.packed_experts: - pf, gs = self.pack_factor, self.group_size + pf, gs = self.quant_config.pack_factor, self.quant_config.group_size self.experts.gate_up_proj_packed.weight_loader = partial( _gate_up_packed_loader, self.tp_rank, self.tp_size, full_inter, ) @@ -303,16 +305,18 @@ def _dispatch_packed_experts( from mstar.utils.fused_moe import fused_experts, moe_sum_reduce_triton reduce = self.tp_size == 1 + # Built per dispatch, not cached: Module._apply rebuilds these Parameters + # on .to(device), so a cached descriptor could hold a stale tensor. + quant = self.quant_config.moe_quant_data( + self.experts.gate_up_proj_scale, self.experts.down_proj_scale, + ) out = fused_experts( flat, self.experts.gate_up_proj_packed, self.experts.down_proj_packed, topk_weights, topk_ids, - w1_scale=self.experts.gate_up_proj_scale, - w2_scale=self.experts.down_proj_scale, - group_size=self.group_size, - pack_factor=self.pack_factor, + quant=quant, reduce_results=reduce, ) if reduce: @@ -329,16 +333,17 @@ def process_weights_after_loading(self, device) -> None: from mstar.model.components.quantization import MarlinMoEMethod from mstar.utils.marlin import is_marlin_available + qc = self.quant_config dev = torch.device(device) shard_inter = divide(self.moe_intermediate_size, self.tp_size) legal_shapes = MarlinMoEMethod.shapes_are_legal( - self.hidden_size, shard_inter, self.group_size + self.hidden_size, shard_inter, qc.group_size ) eligible = ( self.quant_kernel != "triton" and dev.type == "cuda" and torch.cuda.get_device_capability(dev) >= (8, 0) - and self.symmetric + and qc.symmetric and legal_shapes and is_marlin_available() ) @@ -347,7 +352,7 @@ def process_weights_after_loading(self, device) -> None: "quant_kernel='marlin' requested but Marlin is ineligible " f"(needs CUDA sm80+, symmetric INT4, legal shapes: hidden=" f"{self.hidden_size}, shard_inter={shard_inter}, " - f"group_size={self.group_size}, legal={legal_shapes}). " + f"group_size={qc.group_size}, legal={legal_shapes}). " "Use quant_kernel='auto' to fall back to the Triton path." ) global _BACKEND_LOGGED @@ -365,11 +370,11 @@ def process_weights_after_loading(self, device) -> None: logger.info( "KimiSparseMoeBlock routed-expert backend: Marlin W4A16 " "(quant_kernel=%s, group_size=%d, tp_size=%d).", - self.quant_kernel, self.group_size, self.tp_size, + self.quant_kernel, qc.group_size, self.tp_size, ) _BACKEND_LOGGED = True - method = MarlinMoEMethod(num_bits=32 // self.pack_factor, group_size=self.group_size) + method = MarlinMoEMethod.from_quant_config(qc) method.prepare( self.experts.gate_up_proj_packed.data, self.experts.gate_up_proj_scale.data, diff --git a/mstar/model/kimi_k2_7/config.py b/mstar/model/kimi_k2_7/config.py index 9c0afd9b2..19d9410dc 100644 --- a/mstar/model/kimi_k2_7/config.py +++ b/mstar/model/kimi_k2_7/config.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field -from mstar.model.kimi_k2_7.quantization import CompressedTensorsQuantConfig +from mstar.model.components.quantization import CompressedTensorsQuantConfig @dataclass diff --git a/mstar/model/kimi_k2_7/kimi_model.py b/mstar/model/kimi_k2_7/kimi_model.py index 9d6da4883..57d98db0b 100644 --- a/mstar/model/kimi_k2_7/kimi_model.py +++ b/mstar/model/kimi_k2_7/kimi_model.py @@ -48,8 +48,10 @@ def __init__( **kwargs, ): self.cache_dir = cache_dir - checkpoint_path = kwargs.get("checkpoint_path") - self.model_path_hf = checkpoint_path or model_path_hf + # A local checkpoint directory or an HF repo id. Deployments point this at + # their own copy with ``mstar serve --model-path`` / ``mstar-serve + # --model-path`` rather than hardcoding a path in the config yaml. + self.model_path_hf = model_path_hf self._config_variant = kwargs.get("config_variant", "full") if self._config_variant == "reduced": self.config = KimiK2Config.reduced() @@ -327,7 +329,7 @@ def _maybe_apply_checkpoint_quant_config(self, source: str) -> None: import json from pathlib import Path - from mstar.model.kimi_k2_7.quantization import CompressedTensorsQuantConfig + from mstar.model.components.quantization import CompressedTensorsQuantConfig if self.config.quantization_config is not None: return diff --git a/mstar/model/kimi_k2_7/weight_loader.py b/mstar/model/kimi_k2_7/weight_loader.py index eb0ce148f..fa77cc59f 100644 --- a/mstar/model/kimi_k2_7/weight_loader.py +++ b/mstar/model/kimi_k2_7/weight_loader.py @@ -12,7 +12,7 @@ from mstar.model.loader.base import StackedParamRule if TYPE_CHECKING: - from mstar.model.kimi_k2_7.quantization import CompressedTensorsQuantConfig + from mstar.model.components.quantization import CompressedTensorsQuantConfig # Keep the expert index attached while remapping both bf16 and packed sub-keys. _EXPERT_RE = re.compile( @@ -107,7 +107,7 @@ def load_kimi_hf_weights( from mstar.model.loader import load_hf_weights if quant_config is not None: - from mstar.model.kimi_k2_7.quantization import ( + from mstar.model.components.quantization import ( dequant_compressed_tensors_stream, ) diff --git a/mstar/utils/fused_moe/runner.py b/mstar/utils/fused_moe/runner.py index b519d5ea0..2ed68e416 100644 --- a/mstar/utils/fused_moe/runner.py +++ b/mstar/utils/fused_moe/runner.py @@ -19,6 +19,7 @@ invoke_fused_moe_kernel_w4a16, moe_sum_reduce_triton, ) +from mstar.utils.quantization import QuantizationData, QuantizationType def _tl_compute_type(dtype: torch.dtype) -> tl.dtype: @@ -29,6 +30,53 @@ def _tl_compute_type(dtype: torch.dtype) -> tl.dtype: raise ValueError(f"fused_experts: unsupported dtype {dtype}; use bf16 or fp16") +def _validate_expert_shapes( + hidden: int, + w1: torch.Tensor, + w2: torch.Tensor, + quant: QuantizationData | None, +) -> tuple[int, int, int]: + """Check the stacked expert weights against ``hidden`` and the quant scheme. + + Returns ``(num_experts, two_inter, inter)``. Raising here — rather than + dispatching on whichever argument happened to be non-None — is what keeps a + newly-added scheme from silently falling through to the bf16 kernel. + """ + if quant is None: + assert w1.is_contiguous(), "w1 must be contiguous" + assert w2.is_contiguous(), "w2 must be contiguous" + E, two_inter, k_in = w1.shape + assert k_in == hidden, f"w1 last dim {k_in} != hidden {hidden}" + _, w2_hidden, inter = w2.shape + assert w2_hidden == hidden, f"w2 dim[1] {w2_hidden} != hidden {hidden}" + assert two_inter == 2 * inter, f"w1 dim[1] {two_inter} != 2 * w2 dim[2] {2 * inter}" + return E, two_inter, inter + + if quant.quant_type is not QuantizationType.W4A16: + raise ValueError( + f"fused_experts: no kernel for {quant.quant_type}; the Triton MoE path " + f"implements {QuantizationType.W4A16} only" + ) + + pack_factor = quant.pack_factor + assert w1.dtype == torch.int32 and w2.dtype == torch.int32, ( + "W4A16 path expects packed int32 weights" + ) + assert w1.is_contiguous(), "w1 (packed) must be contiguous" + assert w2.is_contiguous(), "w2 (packed) must be contiguous" + E, two_inter, k1_packed = w1.shape + assert k1_packed == hidden // pack_factor, ( + f"w1 packed last dim {k1_packed} != hidden//pack_factor {hidden // pack_factor}" + ) + _, w2_hidden, k2_packed = w2.shape + assert w2_hidden == hidden, f"w2 dim[1] {w2_hidden} != hidden {hidden}" + inter = two_inter // 2 + assert k2_packed == inter // pack_factor, ( + f"w2 packed last dim {k2_packed} != inter//pack_factor {inter // pack_factor}" + ) + return E, two_inter, inter + + def fused_experts( hidden_states: torch.Tensor, w1: torch.Tensor, @@ -37,12 +85,7 @@ def fused_experts( topk_ids: torch.Tensor, activation: str = "silu", reduce_results: bool = True, - w1_scale: torch.Tensor | None = None, - w2_scale: torch.Tensor | None = None, - w1_zp: torch.Tensor | None = None, - w2_zp: torch.Tensor | None = None, - group_size: int | None = None, - pack_factor: int | None = None, + quant: QuantizationData | None = None, ) -> torch.Tensor: """Grouped-GEMM Triton MoE dispatch. @@ -72,12 +115,11 @@ def fused_experts( ``(tokens, hidden)``. If False, skip the sum-reduce and return ``(tokens, top_k, hidden)`` — the caller is responsible for the reduce (e.g. after an all-reduce for TP). - w1_scale, w2_scale : torch.Tensor | None - W4A16 group scales; both ``None`` keeps the historical bf16 path. - w1_zp, w2_zp : torch.Tensor | None - Optional asymmetric zero points (unused for Kimi's symmetric INT4). - group_size, pack_factor : int | None - Required on the W4A16 path. + quant : QuantizationData | None + Scheme descriptor plus its scales/zero points. ``None`` (default) is the + bf16/fp16 path. A :class:`~mstar.utils.quantization.W4A16Data` selects the + packed-INT4 kernel and supplies its group scales; an unimplemented scheme + raises rather than silently taking the bf16 path. Returns ------- @@ -91,36 +133,9 @@ def fused_experts( # Only weights are quantized; activations stay bf16/fp16. assert hidden_states.dtype in (torch.bfloat16, torch.float16) - quantized = w1_scale is not None num_tokens, hidden = hidden_states.shape - if quantized: - assert pack_factor is not None and group_size is not None, ( - "W4A16 path requires pack_factor and group_size" - ) - assert w2_scale is not None, "W4A16 path requires both w1_scale and w2_scale" - assert w1.dtype == torch.int32 and w2.dtype == torch.int32, ( - "W4A16 path expects packed int32 weights" - ) - assert w1.is_contiguous(), "w1 (packed) must be contiguous" - assert w2.is_contiguous(), "w2 (packed) must be contiguous" - E, two_inter, k1_packed = w1.shape - assert k1_packed == hidden // pack_factor, ( - f"w1 packed last dim {k1_packed} != hidden//pack_factor {hidden // pack_factor}" - ) - _, w2_hidden, k2_packed = w2.shape - assert w2_hidden == hidden, f"w2 dim[1] {w2_hidden} != hidden {hidden}" - inter = two_inter // 2 - assert k2_packed == inter // pack_factor, ( - f"w2 packed last dim {k2_packed} != inter//pack_factor {inter // pack_factor}" - ) - else: - assert w1.is_contiguous(), "w1 must be contiguous" - assert w2.is_contiguous(), "w2 must be contiguous" - E, two_inter, k_in = w1.shape - assert k_in == hidden, f"w1 last dim {k_in} != hidden {hidden}" - _, w2_hidden, inter = w2.shape - assert w2_hidden == hidden, f"w2 dim[1] {w2_hidden} != hidden {hidden}" - assert two_inter == 2 * inter, f"w1 dim[1] {two_inter} != 2 * w2 dim[2] {2 * inter}" + E, two_inter, inter = _validate_expert_shapes(hidden, w1, w2, quant) + group_size = quant.group_size if quant is not None else None top_k = topk_ids.shape[1] # moe_align_block_size expects int32; torch.topk returns int64. @@ -154,13 +169,13 @@ def fused_experts( ) # 3. Gate+up GEMM: cache1[slot] = hidden[slot // top_k] @ w1[expert].T - if quantized: + if quant is not None: invoke_fused_moe_kernel_w4a16( A=hidden_states, B_packed=w1, C=cache1, - B_scale=w1_scale, - B_zp=w1_zp, + B_scale=quant.w1_scale, + B_zp=quant.w1_zp, topk_weights=topk_weights, topk_ids=topk_ids, sorted_token_ids=sorted_token_ids, @@ -171,8 +186,8 @@ def fused_experts( config=config, compute_type=compute_type, K=hidden, - pack_factor=pack_factor, - group_size=group_size, + pack_factor=quant.pack_factor, + group_size=quant.group_size, ) else: invoke_fused_moe_kernel( @@ -196,13 +211,13 @@ def fused_experts( # 5. Down GEMM (weighted): cache3[slot] = topk_weight[slot] * (cache2[slot] @ w2[expert].T) # top_k=1 for this GEMM so the kernel's offs_token // top_k is identity # -- it reads cache2 rows directly instead of the (slot // top_k)-th source row. - if quantized: + if quant is not None: invoke_fused_moe_kernel_w4a16( A=cache2, B_packed=w2, C=cache3.view(m_topk, hidden), - B_scale=w2_scale, - B_zp=w2_zp, + B_scale=quant.w2_scale, + B_zp=quant.w2_zp, topk_weights=topk_weights, topk_ids=topk_ids, sorted_token_ids=sorted_token_ids, @@ -213,8 +228,8 @@ def fused_experts( config=config, compute_type=compute_type, K=inter, - pack_factor=pack_factor, - group_size=group_size, + pack_factor=quant.pack_factor, + group_size=quant.group_size, ) else: invoke_fused_moe_kernel( diff --git a/mstar/utils/quantization.py b/mstar/utils/quantization.py new file mode 100644 index 000000000..7fc0435f0 --- /dev/null +++ b/mstar/utils/quantization.py @@ -0,0 +1,90 @@ +"""Quantization scheme descriptors shared by the MoE kernels and model code. + +These are pure data — no ``nn.Module``, no kernel imports — so both +:mod:`mstar.utils.fused_moe` (which consumes them at runtime) and +:mod:`mstar.model.components.quantization` (which produces them) can depend on +this module without a cycle. ``model -> utils`` is the established direction; +defining them model-side would invert it. + +A quantized GEMM entry point takes one :class:`QuantizationData` instead of a +widening list of optional per-scheme kwargs, and branches on +:attr:`QuantizationData.quant_type` rather than sniffing arguments for ``None``. +""" +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from enum import Enum +from typing import ClassVar + +import torch + + +class QuantizationType(Enum): + """Weight/activation precision scheme of a quantized GEMM. + + Members are added alongside a kernel path that implements them — a member + with no dispatch branch is a promise the kernels cannot keep. + """ + + W4A16 = "w4a16" + + +class QuantizationData(ABC): + """Per-call companion data for a quantized GEMM (scales, zero points, layout). + + One subclass per :class:`QuantizationType`. Kernel entry points branch on + :attr:`quant_type` and read the subclass's fields, so supporting a new scheme + is a subclass plus one dispatch branch — not another six optional kwargs whose + legal combinations exist only in a maintainer's head. + + Distinct from a *checkpoint* descriptor such as + ``CompressedTensorsQuantConfig``: that says how the weights were written to + disk (format, ignore list, bit width); this carries the live tensors a + specific kernel call needs. The checkpoint descriptor builds one of these. + """ + + @property + @abstractmethod + def quant_type(self) -> QuantizationType: + """Which scheme this data describes. Kernels dispatch on it.""" + + +@dataclass(frozen=True) +class W4A16Data(QuantizationData): + """Group-wise INT4 weights with bf16/fp16 activations (compressed-tensors layout). + + Weights reach the kernel as low-order-first int32 containers holding + :attr:`pack_factor` nibbles each; ``w1_scale`` / ``w2_scale`` are + ``(num_experts, N, K // group_size)``. ``w1_zp`` / ``w2_zp`` are ``None`` for + symmetric quantization (Kimi-K2.7's case) — that ``None`` *is* the + symmetric/asymmetric discriminant, read off one object rather than re-derived + at each call depth. + + ``frozen=True`` blocks rebinding a field, not mutation of the tensors a field + points at; it marks the object as a description rather than a workspace. + """ + + # The Triton W4A16 kernel extracts nibbles with ``& 0xF`` and subtracts the + # 4-bit offset-binary bias, so the width is fixed by the class, not a knob. + NUM_BITS: ClassVar[int] = 4 + + w1_scale: torch.Tensor + w2_scale: torch.Tensor + group_size: int + w1_zp: torch.Tensor | None = None + w2_zp: torch.Tensor | None = None + + @property + def quant_type(self) -> QuantizationType: + return QuantizationType.W4A16 + + @property + def pack_factor(self) -> int: + """INT4 values per int32 container (8). Derived, so it cannot disagree + with the bit width the kernel actually implements.""" + return 32 // self.NUM_BITS + + @property + def symmetric(self) -> bool: + return self.w1_zp is None and self.w2_zp is None diff --git a/test/integration/test_kimi_mla_absorb_paged.py b/test/integration/test_kimi_mla_absorb_paged.py index 21b8e7f94..483479a7f 100644 --- a/test/integration/test_kimi_mla_absorb_paged.py +++ b/test/integration/test_kimi_mla_absorb_paged.py @@ -4,7 +4,9 @@ from mstar.communication.tensors import LocalTransferEngine from mstar.engine.cache_manager import ( MlaAbsorbCacheManager, + MlaSdpaPlan, WorkspaceBufferManager, + _PlanState, create_cache_manager, ) from mstar.engine.kv_store import ( @@ -137,3 +139,25 @@ def test_paged_latent_mla_real_dims(): def test_paged_latent_mla_reduced_dims(): _run_prefill_then_decode(L=32, Drope=8, H=4, T=6, page_size=4) + + +def test_sdpa_plan_clears_any_injected_wrapper(): + """``run_attention_mla`` picks its path with ``ps.wrapper is not None``, so the + SDPA branch must clear the slot — otherwise a wrapper injected by CUDA-graph + capture routes latent attention into a paged kernel that cannot read the 4-D + cache. Mirrors how the paged path clears ``dense_gen``.""" + L, Drope = 32, 8 # reduced dims: no MLA kernel, so plan_attention takes SDPA + cm, alloc = _make_latent_cache_manager(L + Drope, torch.bfloat16, 0.1) + try: + cm.set_active_label("main") + # Stand in for a persistent wrapper injected via cuda_graph_plan_states. + cm._plan_states["main"] = _PlanState(wrapper=object()) + + cm.plan_attention(seq_lens=[4], is_causal=True, dtype=torch.bfloat16) + + ps = cm._plan_states["main"] + assert ps.wrapper is None, "SDPA plan left a stale wrapper in the plan state" + assert isinstance(ps.mla, MlaSdpaPlan) + assert [r.seq_len for r in ps.mla.requests] == [4] + finally: + alloc.cleanup() diff --git a/test/integration/test_kimi_moe_inkernel_dequant.py b/test/integration/test_kimi_moe_inkernel_dequant.py index 062813481..7e4a246f5 100644 --- a/test/integration/test_kimi_moe_inkernel_dequant.py +++ b/test/integration/test_kimi_moe_inkernel_dequant.py @@ -1,8 +1,8 @@ import pytest import torch +from mstar.model.components.quantization import W4A16Data, unpack_int32 from mstar.model.kimi_k2_7._testing import fake_quantize_weight -from mstar.model.kimi_k2_7.quantization import unpack_int32 pytestmark = pytest.mark.skipif( not torch.cuda.is_available(), @@ -58,7 +58,7 @@ def test_w4a16_matches_bf16_on_same_dequant(num_tokens): out_quant = fused_experts( x, w1_packed, w2_packed, topk_weights, topk_ids, - w1_scale=w1_scale, w2_scale=w2_scale, group_size=GROUP_SIZE, pack_factor=PACK_FACTOR, + quant=W4A16Data(w1_scale=w1_scale, w2_scale=w2_scale, group_size=GROUP_SIZE), ) out_bf16 = fused_experts(x, w1_deq, w2_deq, topk_weights, topk_ids) @@ -81,7 +81,7 @@ def test_w4a16_reduce_results_false_shape(): got = fused_experts( x, w1_packed, w2_packed, topk_weights, topk_ids, - w1_scale=w1_scale, w2_scale=w2_scale, group_size=GROUP_SIZE, pack_factor=PACK_FACTOR, + quant=W4A16Data(w1_scale=w1_scale, w2_scale=w2_scale, group_size=GROUP_SIZE), reduce_results=False, ) exp = fused_experts(x, w1_deq, w2_deq, topk_weights, topk_ids, reduce_results=False) diff --git a/test/integration/test_kimi_serve_e2e.py b/test/integration/test_kimi_serve_e2e.py index 537dd867a..f0bd9edbd 100644 --- a/test/integration/test_kimi_serve_e2e.py +++ b/test/integration/test_kimi_serve_e2e.py @@ -198,8 +198,8 @@ def test_serve_path_prefill_decode_loop(tmp_path): cfg = _write_checkpoint(tmp_path, seed=0) model = KimiK2Model( - model_path_hf="", config_variant="reduced", - checkpoint_path=str(tmp_path), tokenizer_mode="byte", + model_path_hf=str(tmp_path), config_variant="reduced", + tokenizer_mode="byte", ) assert model.config.vocab_size == 256 @@ -228,8 +228,8 @@ def test_serve_path_prefill_decode_loop(tmp_path): def test_serve_path_is_deterministic(tmp_path): cfg = _write_checkpoint(tmp_path, seed=1) model = KimiK2Model( - model_path_hf="", config_variant="reduced", - checkpoint_path=str(tmp_path), tokenizer_mode="byte", + model_path_hf=str(tmp_path), config_variant="reduced", + tokenizer_mode="byte", ) submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16) prompt_ids = model.process_prompt("serve", ["text"], ["text"])["text_inputs"][0].to(DEVICE) diff --git a/test/integration/test_marlin_kernels.py b/test/integration/test_marlin_kernels.py index 8b012d1c4..764ac1b9b 100644 --- a/test/integration/test_marlin_kernels.py +++ b/test/integration/test_marlin_kernels.py @@ -2,6 +2,7 @@ import pytest import torch +from mstar.model.components.quantization import W4A16Data from mstar.model.kimi_k2_7._testing import fake_quantize_weight pytestmark = pytest.mark.skipif( @@ -83,7 +84,7 @@ def test_marlin_moe_matches_bf16_and_triton(num_tokens): out_bf16 = fused_experts(x, w1_deq, w2_deq, topk_weights, topk_ids) # ground truth out_triton = fused_experts( # same packed nibbles, Triton W4A16 kernel x, w1_packed, w2_packed, topk_weights, topk_ids, - w1_scale=w1_scale, w2_scale=w2_scale, group_size=GROUP_SIZE, pack_factor=PACK_FACTOR, + quant=W4A16Data(w1_scale=w1_scale, w2_scale=w2_scale, group_size=GROUP_SIZE), ) assert out_marlin.shape == (num_tokens, H) and out_marlin.dtype == torch.bfloat16 diff --git a/test/modular/test_cache_manager_backends.py b/test/modular/test_cache_manager_backends.py index 36039b418..99d962030 100644 --- a/test/modular/test_cache_manager_backends.py +++ b/test/modular/test_cache_manager_backends.py @@ -1,8 +1,13 @@ """KV-cache attention-backend selection: ``KVCacheConfig.attention_backend`` names the ``BatchedCacheManager`` subclass ``create_cache_manager`` -instantiates, unknown names are rejected, and the base class is abstract.""" +instantiates, unknown names are rejected, and the base class is abstract. + +Also covers the ``mla_absorb`` kernel predicate and the CUDA-graph capture +decision it drives — the two must agree, or capture builds a paged wrapper that +cannot read the 4-D latent cache.""" import pytest +import torch from mstar.engine.cache_manager import ( ATTENTION_BACKENDS, @@ -10,7 +15,9 @@ DenseGenCacheManager, FlashInferCacheManager, create_cache_manager, + mla_kernel_available_for, ) +from mstar.engine.cuda_graph_runner import mla_absorb_capture_blocked from mstar.engine.kv_store import KVCacheConfig @@ -78,3 +85,79 @@ def test_base_class_is_abstract(): def test_registry_names(): assert set(ATTENTION_BACKENDS) == {"flashinfer", "dense_gen", "mla_absorb"} + + +# --- mla_absorb: kernel predicate + capture decision ---------------------- + +# Real Kimi latent dims; flashinfer's MLA wrapper is hard-locked to these. +_REAL_CKV, _REAL_KPE = 512, 64 +_CUDA = torch.device("cuda:0") # constructible without a GPU present + + +def _mla_cfg(mla_ckv_dim: int | None, head_dim: int) -> KVCacheConfig: + return KVCacheConfig( + num_layers=2, num_kv_heads=1, head_dim=head_dim, max_seq_len=64, + attention_backend="mla_absorb", mla_ckv_dim=mla_ckv_dim, + ) + + +@pytest.fixture +def sm90(monkeypatch): + """Pretend we are on a Hopper GPU, without needing one.""" + import mstar.engine.cache_manager as cm_mod + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda d: (9, 0)) + # _mla_kernel_available is functools.cache'd and imports flashinfer; stub the + # module attribute so the predicate's own logic is what's under test. + monkeypatch.setattr( + cm_mod, "_mla_kernel_available", + lambda ckv, kpe, sm: (ckv, kpe, sm) == (_REAL_CKV, _REAL_KPE, 9), + ) + + +def test_kernel_predicate_true_only_for_real_dims_on_sm90(sm90): + cfg = _mla_cfg(_REAL_CKV, _REAL_CKV + _REAL_KPE) + assert mla_kernel_available_for(cfg, _CUDA) is True + + +def test_kernel_predicate_false_for_reduced_dims(sm90): + # The reduced test config: right backend, wrong latent width. + assert mla_kernel_available_for(_mla_cfg(32, 40), _CUDA) is False + + +def test_kernel_predicate_false_without_ckv_dim(sm90): + assert mla_kernel_available_for(_mla_cfg(None, 40), _CUDA) is False + + +def test_kernel_predicate_is_total_on_cpu(sm90): + """Must return False, not raise: torch.cuda.get_device_capability rejects a + CPU device, and the absorbed-SDPA fallback has to stay reachable there.""" + cfg = _mla_cfg(_REAL_CKV, _REAL_CKV + _REAL_KPE) + assert mla_kernel_available_for(cfg, torch.device("cpu")) is False + + +def test_capture_not_blocked_for_other_backends(sm90): + for backend in ("flashinfer", "dense_gen"): + assert mla_absorb_capture_blocked(_cfg(backend), _CUDA) is None + + +def test_capture_not_blocked_when_kernel_serves_the_dims(sm90): + cfg = _mla_cfg(_REAL_CKV, _REAL_CKV + _REAL_KPE) + assert mla_absorb_capture_blocked(cfg, _CUDA) is None + + +@pytest.mark.parametrize( + "ckv, head_dim", [(32, 40), (None, 40), (_REAL_CKV, _REAL_CKV + 128)] +) +def test_capture_blocked_when_absorbed_falls_back_to_sdpa(sm90, ckv, head_dim): + """No wrapper is planned on the SDPA path, so a captured graph could only + hold a paged wrapper — capture must be cancelled, not attempted.""" + reason = mla_absorb_capture_blocked(_mla_cfg(ckv, head_dim), _CUDA) + assert reason is not None + assert "mla_absorb" in reason and "eager" in reason + + +def test_capture_blocked_on_cpu_device(sm90): + cfg = _mla_cfg(_REAL_CKV, _REAL_CKV + _REAL_KPE) + assert mla_absorb_capture_blocked(cfg, torch.device("cpu")) is not None diff --git a/test/modular/test_kimi_quant.py b/test/modular/test_compressed_tensors.py similarity index 99% rename from test/modular/test_kimi_quant.py rename to test/modular/test_compressed_tensors.py index 390a649db..1e9a7c098 100644 --- a/test/modular/test_kimi_quant.py +++ b/test/modular/test_compressed_tensors.py @@ -1,14 +1,14 @@ import pytest import torch -from mstar.model.kimi_k2_7._testing import fake_quantize_weight -from mstar.model.kimi_k2_7.quantization import ( +from mstar.model.components.quantization import ( CompressedTensorsQuantConfig, dequant_compressed_tensors_stream, dequantize_weight, pack_int32, unpack_int32, ) +from mstar.model.kimi_k2_7._testing import fake_quantize_weight def test_pack_known_answer(): diff --git a/test/modular/test_quantization_types.py b/test/modular/test_quantization_types.py new file mode 100644 index 000000000..a90840179 --- /dev/null +++ b/test/modular/test_quantization_types.py @@ -0,0 +1,136 @@ +"""Quantization descriptors: the tagged union that replaced ``fused_experts``' +per-scheme kwargs, and the checkpoint-config -> kernel-data seam.""" + +import dataclasses + +import pytest +import torch + +from mstar.model.components.quantization import ( + CompressedTensorsQuantConfig, + MarlinMoEMethod, + QuantizationData, + QuantizationType, + W4A16Data, +) + + +def _data(**overrides) -> W4A16Data: + kwargs = { + "w1_scale": torch.ones(2, 4, 2), + "w2_scale": torch.ones(2, 4, 2), + "group_size": 32, + } + kwargs.update(overrides) + return W4A16Data(**kwargs) + + +# --- the union ------------------------------------------------------------ + + +def test_quantization_data_is_abstract(): + with pytest.raises(TypeError, match="abstract"): + QuantizationData() + + +def test_w4a16_tags_itself_and_derives_pack_factor(): + d = _data() + assert d.quant_type is QuantizationType.W4A16 + # Derived from the class's NUM_BITS, so it cannot disagree with the kernel's + # hardcoded 4-bit nibble extraction. + assert d.pack_factor == 8 + assert "pack_factor" not in {f.name for f in dataclasses.fields(d)} + assert "NUM_BITS" not in {f.name for f in dataclasses.fields(d)} + + +def test_zero_points_are_the_symmetry_discriminant(): + assert _data().symmetric is True + assert _data(w1_zp=torch.zeros(2, 4, 2), w2_zp=torch.zeros(2, 4, 2)).symmetric is False + + +def test_frozen_blocks_rebinding_but_not_tensor_mutation(): + d = _data() + with pytest.raises(dataclasses.FrozenInstanceError): + d.group_size = 64 + # Documented limit: frozen freezes the binding, not the pointee. + d.w1_scale[0, 0, 0] = 7.0 + assert d.w1_scale[0, 0, 0] == 7.0 + + +# --- the checkpoint-config -> kernel-data seam ---------------------------- + + +def test_config_maps_num_bits_to_scheme(): + assert CompressedTensorsQuantConfig(num_bits=4).quant_type is QuantizationType.W4A16 + + +@pytest.mark.parametrize("num_bits", [2, 8, 16]) +def test_config_rejects_unimplemented_width_at_load(num_bits): + # The whole point: an INT8 checkpoint fails here, naming the checkpoint, + # instead of reaching a kernel that would mask its 8-bit values to nibbles. + cfg = CompressedTensorsQuantConfig(num_bits=num_bits) + assert cfg.quant_type is None # the query answers; only the ensure raises + with pytest.raises(ValueError, match="mstar does not implement"): + cfg.ensure_kernel_support() + + +def test_ensure_kernel_support_returns_the_scheme(): + cfg = CompressedTensorsQuantConfig(num_bits=4) + assert cfg.ensure_kernel_support() is QuantizationType.W4A16 + + +def test_moe_quant_data_round_trips_group_size_and_scales(): + cfg = CompressedTensorsQuantConfig(num_bits=4, group_size=64) + w1s, w2s = torch.ones(2, 4, 2), torch.zeros(2, 4, 2) + data = cfg.moe_quant_data(w1s, w2s) + + assert isinstance(data, W4A16Data) + assert data.quant_type is QuantizationType.W4A16 + assert data.group_size == 64 + assert data.pack_factor == cfg.pack_factor + assert data.w1_scale is w1s and data.w2_scale is w2s + + +def test_symmetric_config_drops_zero_points(): + sym = CompressedTensorsQuantConfig(num_bits=4, symmetric=True) + zp = torch.zeros(2, 4, 2) + data = sym.moe_quant_data(torch.ones(2, 4, 2), torch.ones(2, 4, 2), w1_zp=zp, w2_zp=zp) + assert data.w1_zp is None and data.w2_zp is None and data.symmetric + + asym = CompressedTensorsQuantConfig(num_bits=4, symmetric=False) + data = asym.moe_quant_data(torch.ones(2, 4, 2), torch.ones(2, 4, 2), w1_zp=zp, w2_zp=zp) + assert data.w1_zp is zp and not data.symmetric + + +def test_marlin_reads_num_bits_forward_from_the_config(): + cfg = CompressedTensorsQuantConfig(num_bits=4, group_size=128) + method = MarlinMoEMethod.from_quant_config(cfg) + assert method.num_bits == 4 + assert method.group_size == 128 + assert method.quant_type is QuantizationType.W4A16 + + with pytest.raises(ValueError, match="mstar does not implement"): + MarlinMoEMethod.from_quant_config(CompressedTensorsQuantConfig(num_bits=8)) + + +# --- kernel dispatch ------------------------------------------------------ + + +def test_fused_experts_rejects_an_unhandled_scheme(): + """A future QuantizationType with no branch must raise, not silently take + the bf16 path — that is what the sentinel-based dispatch could not do.""" + pytest.importorskip("triton") + from mstar.utils.fused_moe.runner import _validate_expert_shapes + + class _FutureData(QuantizationData): + @property + def quant_type(self): + return "int8-someday" + + with pytest.raises(ValueError, match="no kernel for"): + _validate_expert_shapes( + hidden=128, + w1=torch.zeros(2, 8, 16, dtype=torch.int32), + w2=torch.zeros(2, 128, 1, dtype=torch.int32), + quant=_FutureData(), + )