Ran linter

Author: Rohan-Bierneni

src/maxtext/integration/vllm/torchax_converter/qwen35_moe.py

@@ -28,10 +28,8 @@ class Qwen35MaxTextToVLLMConverter(BaseMaxTextToVLLMConverter):
   NUM_SLOTS = 4  # 3 GDN layers + 1 Full Attention layer per cycle
 
   def convert(self, model_state: dict):
-    """Main entry point for the Tunix weight synchronization."""
     logging.info("\n%sStarting Qwen 3.5 Conversion (Hybrid 3:1 MoE)...%s", GREEN, RESET)
     self.vllm_state = {}
-
     self.num_reps = self.num_layers // self.NUM_SLOTS
 
     with timer("Convert Global Weights"):
@@ -43,45 +41,27 @@ def convert(self, model_state: dict):
     with timer("Convert MoE Weights"):
       self._convert_moe(model_state)
 
-    # ------------------------------------------------------------------ #
-    # Protect JAX compilation
-    # ------------------------------------------------------------------ #
+    # Protect JAX compilation by enforcing bfloat16
     for key in self.vllm_state:
       self.vllm_state[key] = self.vllm_state[key].astype(jnp.bfloat16)
 
     return self.vllm_state
 
-  # ------------------------------------------------------------------ #
-  # 1. Global Weights
-  # ------------------------------------------------------------------ #
   def _convert_global(self, params):
-    logging.info("_convert_global: Processing embeddings and LM head...")
-
     self.vllm_state["vllm_model.language_model.model.embed_tokens.weight"] = jnp.array(
         params["base"]["token_embedder"]["embedding"]
     )
-
     self.vllm_state["vllm_model.language_model.model.norm.weight"] = jnp.array(
         params["base"]["decoder"]["decoder_norm"]["scale"]
     )
-
     self.vllm_state["vllm_model.language_model.lm_head.weight"] = jnp.transpose(
         params["base"]["decoder"]["logits_dense"]["kernel"], (1, 0)
     )
 
-  # ------------------------------------------------------------------ #
-  # 2. Hybrid Attention (Scanned 3:1 Blocks)
-  # ------------------------------------------------------------------ #
   def _convert_attn(self, params):
-    logging.info("_convert_attn: Unstacking layer norms and routing hybrid attention...")
     decoder = params["base"]["decoder"]
-
-    if "scanned_blocks" in decoder:
-      blocks = decoder["scanned_blocks"]
-      slot_prefix = "layers"
-    else:
-      blocks = decoder["layers"]
-      slot_prefix = "layer"
+    blocks = decoder.get("scanned_blocks", decoder.get("layers"))
+    slot_prefix = "layers" if "scanned_blocks" in decoder else "layer"
 
     @jax.jit
     def _unstack_rep(x):
@@ -114,20 +94,18 @@ def _unstack_rep(x):
 
           q, k, v = q_layers[rep], k_layers[rep], v_layers[rep]
 
-          # Transpose to standard (num_heads, head_dim, emb_dim)
           q_T = jnp.transpose(q, (1, 2, 0))
           k_T = jnp.transpose(k, (1, 2, 0))
           v_T = jnp.transpose(v, (1, 2, 0))
 
-          # Flatten head dimensions and slice for TP interleaving
           tp_size = self.vllm_tp
           q_tp_shards = jnp.split(q_T.reshape(-1, q.shape[0]), tp_size, axis=0)
           k_tp_shards = jnp.split(k_T.reshape(-1, k.shape[0]), tp_size, axis=0)
           v_tp_shards = jnp.split(v_T.reshape(-1, v.shape[0]), tp_size, axis=0)
 
-          tp_interleaved = []
-          for t in range(tp_size):
-            tp_interleaved.append(jnp.concatenate([q_tp_shards[t], k_tp_shards[t], v_tp_shards[t]], axis=0))
+          tp_interleaved = [
+              jnp.concatenate([q_tp_shards[t], k_tp_shards[t], v_tp_shards[t]], axis=0) for t in range(tp_size)
+          ]
 
           self.vllm_state[f"{prefix}.self_attn.qkv_proj.weight"] = jnp.concatenate(tp_interleaved, axis=0)
           self.vllm_state[f"{prefix}.self_attn.o_proj.weight"] = jnp.transpose(o_layers[rep], (1, 0))
@@ -136,11 +114,9 @@ def _unstack_rep(x):
 
       else:
         gdn = slot_data["attention"]
-
         qkvz_layers = jnp.unstack(gdn["in_proj_qkvz"]["kernel"], axis=1)
         ba_layers = jnp.unstack(gdn["in_proj_ba"]["kernel"], axis=1)
         out_layers = jnp.unstack(gdn["out_proj"]["kernel"], axis=1)
-
         conv_layers = jnp.unstack(gdn["conv1d"]["kernel"], axis=1)
 
         A_log_layers = jnp.unstack(gdn["A_log"], axis=1)
@@ -154,84 +130,55 @@ def _unstack_rep(x):
           self.vllm_state[f"{prefix}.input_layernorm.weight"] = pre_ln[rep]
           self.vllm_state[f"{prefix}.post_attention_layernorm.weight"] = post_ln[rep]
 
-          # Extract MaxText QKVZ layout
-          H_k = 16
-          H_v = 32
-          D_k = 128
-          D_v = 128
-          V_per_K = 2
+          # Extract MaxText GDN QKVZ Layout
+          H_k, H_v, D_k, D_v, V_per_K = 16, 32, 128, 128, 2
 
           t_m = jnp.transpose(qkvz_layers[rep], (1, 0))
           block_size = D_k + D_k + V_per_K * D_v + V_per_K * D_v
           t_r = t_m.reshape(H_k, block_size, -1)
 
-          q_r = t_r[:, :D_k, :]
-          k_r = t_r[:, D_k : 2 * D_k, :]
-          v_r = t_r[:, 2 * D_k : 2 * D_k + V_per_K * D_v, :]
-          z_r = t_r[:, 2 * D_k + V_per_K * D_v :, :]
-
-          q = q_r.reshape(H_k * D_k, -1)
-          k = k_r.reshape(H_k * D_k, -1)
-          v = v_r.reshape(H_v * D_v, -1)
-          z = z_r.reshape(H_v * D_v, -1)
+          q = t_r[:, :D_k, :].reshape(H_k * D_k, -1)
+          k = t_r[:, D_k : 2 * D_k, :].reshape(H_k * D_k, -1)
+          v = t_r[:, 2 * D_k : 2 * D_k + V_per_K * D_v, :].reshape(H_v * D_v, -1)
+          z = t_r[:, 2 * D_k + V_per_K * D_v :, :].reshape(H_v * D_v, -1)
 
-          # Interleave GDN QKVZ by Tensor Parallel shard
           tp_size = self.vllm_tp
           q_shards = jnp.split(q, tp_size, axis=0)
           k_shards = jnp.split(k, tp_size, axis=0)
           v_shards = jnp.split(v, tp_size, axis=0)
           z_shards = jnp.split(z, tp_size, axis=0)
 
-          qkvz_interleaved_shards = []
-          for s in range(tp_size):
-            qkvz_interleaved_shards.append(jnp.concatenate([q_shards[s], k_shards[s], v_shards[s], z_shards[s]], axis=0))
+          qkvz_interleaved = [
+              jnp.concatenate([q_shards[s], k_shards[s], v_shards[s], z_shards[s]], axis=0) for s in range(tp_size)
+          ]
+          self.vllm_state[f"{prefix}.linear_attn.in_proj_qkvz.weight"] = jnp.concatenate(qkvz_interleaved, axis=0)
 
-          self.vllm_state[f"{prefix}.linear_attn.in_proj_qkvz.weight"] = jnp.concatenate(qkvz_interleaved_shards, axis=0)
-
-          # Extract MaxText BA layout
+          # Extract MaxText GDN BA Layout
           t_m_ba = jnp.transpose(ba_layers[rep], (1, 0))
           block_size_ba = V_per_K * 2
           t_r_ba = t_m_ba.reshape(H_k, block_size_ba, -1)
 
-          b_r = t_r_ba[:, :V_per_K, :]
-          a_r = t_r_ba[:, V_per_K:, :]
-
-          b = b_r.reshape(H_v, -1)
-          a = a_r.reshape(H_v, -1)
+          b = t_r_ba[:, :V_per_K, :].reshape(H_v, -1)
+          a = t_r_ba[:, V_per_K:, :].reshape(H_v, -1)
 
-          # Interleave BA vectors by Tensor Parallel shard
           b_shards = jnp.split(b, tp_size, axis=0)
           a_shards = jnp.split(a, tp_size, axis=0)
 
-          ba_interleaved_shards = []
-          for s in range(tp_size):
-            ba_interleaved_shards.append(jnp.concatenate([b_shards[s], a_shards[s]], axis=0))
+          ba_interleaved = [jnp.concatenate([b_shards[s], a_shards[s]], axis=0) for s in range(tp_size)]
+          self.vllm_state[f"{prefix}.linear_attn.in_proj_ba.weight"] = jnp.concatenate(ba_interleaved, axis=0)
 
-          self.vllm_state[f"{prefix}.linear_attn.in_proj_ba.weight"] = jnp.concatenate(ba_interleaved_shards, axis=0)
           self.vllm_state[f"{prefix}.linear_attn.out_proj.weight"] = jnp.transpose(out_layers[rep], (1, 0))
-
-          # MT: [K, 1, C] <-> HF: [C, 1, K]
-          conv_w = conv_layers[rep]
-          self.vllm_state[f"{prefix}.linear_attn.conv1d.weight"] = jnp.transpose(conv_w, (2, 1, 0))
+          self.vllm_state[f"{prefix}.linear_attn.conv1d.weight"] = jnp.transpose(conv_layers[rep], (2, 1, 0))
           self.vllm_state[f"{prefix}.linear_attn.A_log"] = A_log_layers[rep]
           self.vllm_state[f"{prefix}.linear_attn.dt_bias"] = dt_bias_layers[rep]
           self.vllm_state[f"{prefix}.linear_attn.norm.weight"] = gdn_norm_layers[rep]
 
       gc.collect()
 
-  # ------------------------------------------------------------------ #
-  # 3. Mixture of Experts (Scanned Block)
-  # ------------------------------------------------------------------ #
   def _convert_moe(self, params):
-    logging.info("_convert_moe: Packaging routed and shared experts...")
     decoder = params["base"]["decoder"]
-
-    if "scanned_blocks" in decoder:
-      blocks = decoder["scanned_blocks"]
-      slot_prefix = "layers"
-    else:
-      blocks = decoder["layers"]
-      slot_prefix = "layer"
+    blocks = decoder.get("scanned_blocks", decoder.get("layers"))
+    slot_prefix = "layers" if "scanned_blocks" in decoder else "layer"
 
     for slot in range(self.NUM_SLOTS):
       slot_data = blocks[f"{slot_prefix}_{slot}"]
@@ -245,21 +192,35 @@ def _convert_moe(self, params):
 
       router_weights = jnp.unstack(jnp.transpose(routed["gate"]["kernel"], (1, 2, 0)), axis=0)
 
-      # Fusing and Tensor Parallel Interleaving for MoE W1 and W3
+      # -------------------------------------------------------------
+      # Fusing, TP Interleaving, and TPU GMM Alignment for W1 and W3
+      # -------------------------------------------------------------
       wi_0 = jnp.transpose(routed["wi_0"], (1, 0, 2, 3))
       wi_1 = jnp.transpose(routed["wi_1"], (1, 0, 2, 3))
 
+      num_reps, num_experts, d_model, d_inner = wi_0.shape
       tp_size = self.vllm_tp
-      w1_shards = jnp.split(wi_0, tp_size, axis=-1)
-      w3_shards = jnp.split(wi_1, tp_size, axis=-1)
 
-      interleaved_shards = []
-      for i in range(tp_size):
-        interleaved_shards.append(w1_shards[i])
-        interleaved_shards.append(w3_shards[i])
+      # vLLM's TPU Grouped GEMM kernel requires 128-alignment per expert chunk
+      chunk_size = d_inner // tp_size
+      padded_chunk_size = ((chunk_size + 127) // 128) * 128
+      pad_amount = padded_chunk_size - chunk_size
+
+      w1_chunks = wi_0.reshape(num_reps, num_experts, d_model, tp_size, chunk_size)
+      w3_chunks = wi_1.reshape(num_reps, num_experts, d_model, tp_size, chunk_size)
+
+      # Apply padding if running on a topology that splinters chunks below 128 (e.g. TP=8)
+      if pad_amount > 0:
+        w1_chunks = jnp.pad(w1_chunks, ((0, 0), (0, 0), (0, 0), (0, 0), (0, pad_amount)))
+        w3_chunks = jnp.pad(w3_chunks, ((0, 0), (0, 0), (0, 0), (0, 0), (0, pad_amount)))
+
+      # Interleave W1 and W3 shards -> Shape: (reps, exp, d_model, tp, 2, padded_chunk)
+      combined_shards = jnp.stack([w1_chunks, w3_chunks], axis=-2)
 
-      gate_up = jnp.concatenate(interleaved_shards, axis=-1)
+      # Flatten the TP, 2, and chunk dimensions back into the final inner dimension
+      gate_up = combined_shards.reshape(num_reps, num_experts, d_model, -1)
       w13_layers = jnp.unstack(gate_up, axis=0)
+      # -------------------------------------------------------------
 
       wo_transposed = jnp.transpose(routed["wo"], (1, 0, 2, 3))
       down_layers = jnp.unstack(wo_transposed, axis=0)
@@ -282,7 +243,6 @@ def _convert_moe(self, params):
         self.vllm_state[f"{p}.mlp.experts.w13_weight"] = w13_layers[rep]
         self.vllm_state[f"{p}.mlp.experts.w2_weight"] = down_layers[rep]
 
-        # Build Shared Expert structure
         if has_shared:
           sh_g, sh_u = sh_gate_layers[rep], sh_up_layers[rep]
           sh_per_tp = sh_g.shape[0] // self.vllm_tp