fix

Author: cjx0709

python/sgl_jax/srt/layers/attention/fla/group_rmsnorm.py

@@ -30,6 +30,14 @@ def __init__(
     ):
         if hidden_size % num_groups != 0:
             raise ValueError("hidden_size must be divisible by num_groups")
+        if mesh is not None:
+            tp_size = mesh.shape.get("tensor", 1)
+            if tp_size < num_groups:
+                raise ValueError(
+                    "GroupRMSNorm requires tensor parallel size to be at least "
+                    f"num_groups to keep each RMS group intact, got tensor "
+                    f"parallel size={tp_size}, num_groups={num_groups}."
+                )
 
         self.hidden_size = hidden_size
         self.num_groups = num_groups

python/sgl_jax/srt/layers/attention/linear/lightning_backend.py

@@ -16,7 +16,6 @@
 
 import logging
 import math
-import os
 from typing import TYPE_CHECKING
 
 import jax
@@ -32,19 +31,13 @@
 logger = logging.getLogger(__name__)
 
 try:
-    from sgl_jax.srt.kernels.simple_gla.native import (
-        naive_gla_decode,
-        naive_gla_prefill,
-    )
     from sgl_jax.srt.kernels.simple_gla.simple_gla import (
         fused_recurrent_simple_gla,
         simple_gla_fwd,
     )
 except ModuleNotFoundError:
     simple_gla_fwd = None
     fused_recurrent_simple_gla = None
-    naive_gla_decode = None
-    naive_gla_prefill = None
 
 if TYPE_CHECKING:
     from sgl_jax.srt.layers.radix_lightning_attention import RadixLightningAttention
@@ -107,7 +100,6 @@ def __init__(
         """
         super().__init__(mesh=mesh)
         self.chunk_size = chunk_size
-        self.use_native_gla = os.environ.get("SGLANG_JAX_GLA_BACKEND", "").lower() == "native"
         if (
             linear_recurrent_layer_ids is not None
             and num_hidden_layers is not None
@@ -212,7 +204,7 @@ def _forward_decode(
         slope: jnp.ndarray,
     ) -> tuple[jax.Array, jax.Array]:
         """Decode forward using shard_map."""
-        if fused_recurrent_simple_gla is None or naive_gla_decode is None:
+        if fused_recurrent_simple_gla is None:
             raise ImportError("simple_gla kernel is required for GLA decode")
 
         ssm_states = ssm_states.astype(jnp.float32)
@@ -221,25 +213,15 @@ def _decode_fn(q_local, k_local, v_local, gamma, h0):
             q_d = q_local[:, None, :, :]
             k_d = k_local[:, None, :, :]
             v_d = v_local[:, None, :, :]
-            if self.use_native_gla:
-                output_d, new_state = naive_gla_decode(
-                    q_d,
-                    k_d,
-                    v_d,
-                    g_gamma=gamma,
-                    h0=h0,
-                    scale=None,
-                )
-            else:
-                output_d, new_state = fused_recurrent_simple_gla(
-                    q_d,
-                    k_d,
-                    v_d,
-                    g_gamma=gamma,
-                    initial_state=h0,
-                    output_final_state=True,
-                    scale=None,
-                )
+            output_d, new_state = fused_recurrent_simple_gla(
+                q_d,
+                k_d,
+                v_d,
+                g_gamma=gamma,
+                initial_state=h0,
+                output_final_state=True,
+                scale=None,
+            )
             return output_d[:, 0, :, :], new_state
 
         output, new_state = jax.shard_map(
@@ -270,36 +252,25 @@ def _forward_extend(
         slope: jnp.ndarray,
     ) -> tuple[jax.Array, jax.Array]:
         """Extend forward using shard_map."""
-        if simple_gla_fwd is None or naive_gla_prefill is None:
+        if simple_gla_fwd is None:
             raise ImportError("simple_gla kernel is required for GLA prefill")
 
         cu_seqlens = self.forward_metadata.cu_q_lens
         ssm_states = ssm_states.astype(jnp.float32)
         chunk_size = self.chunk_size
 
         def _prefill_fn(q_local, k_local, v_local, gamma, h0, cu_seqlens_p):
-            if self.use_native_gla:
-                output, ht = naive_gla_prefill(
-                    q_local[None],
-                    k_local[None],
-                    v_local[None],
-                    g_gamma=gamma,
-                    h0=h0,
-                    cu_seqlens=cu_seqlens_p,
-                    scale=None,
-                )
-            else:
-                output, ht = simple_gla_fwd(
-                    q_local[None],
-                    k_local[None],
-                    v_local[None],
-                    g_gamma=gamma,
-                    h0=h0,
-                    cu_seqlens_dev=cu_seqlens_p,
-                    scale=None,
-                    use_ht=True,
-                    chunk_size=chunk_size,
-                )
+            output, ht = simple_gla_fwd(
+                q_local[None],
+                k_local[None],
+                v_local[None],
+                g_gamma=gamma,
+                h0=h0,
+                cu_seqlens_dev=cu_seqlens_p,
+                scale=None,
+                use_ht=True,
+                chunk_size=chunk_size,
+            )
             return output[0], ht
 
         output, new_state = jax.shard_map(

python/sgl_jax/test/layers/test_group_rmsnorm.py

@@ -1,5 +1,9 @@
+import jax
 import jax.numpy as jnp
 import numpy as np
+import pytest
+from jax.sharding import AxisType, Mesh, NamedSharding
+from jax.sharding import PartitionSpec as P
 
 from sgl_jax.srt.layers.attention.fla.group_rmsnorm import GroupRMSNorm
 
@@ -39,17 +43,42 @@ def _make_weight(rng, hidden_size=HIDDEN_SIZE):
     return rng.standard_normal(hidden_size).astype(np.float32)
 
 
+def _make_mesh(num_groups=NUM_GROUPS):
+    devices = np.array(jax.devices())
+    if devices.size < num_groups:
+        pytest.skip(
+            f"GroupRMSNorm sharded test requires at least {num_groups} devices, got {devices.size}"
+        )
+    return Mesh(
+        devices[:num_groups].reshape(1, num_groups),
+        axis_names=("data", "tensor"),
+        axis_types=(AxisType.Explicit, AxisType.Explicit),
+    )
+
+
 def _make_jax_model(hidden_size=HIDDEN_SIZE, num_groups=NUM_GROUPS, weight=None):
     """Create a JAX GroupRMSNorm model, optionally with custom weight."""
-    model = GroupRMSNorm(hidden_size, num_groups=num_groups, epsilon=EPSILON)
+    mesh = _make_mesh(num_groups)
+    with jax.set_mesh(mesh):
+        model = GroupRMSNorm(
+            hidden_size,
+            num_groups=num_groups,
+            epsilon=EPSILON,
+            kernel_axes=("tensor",),
+            mesh=mesh,
+        )
     if weight is not None:
-        model.weight[...] = jnp.array(weight)
+        model.weight[...] = jax.device_put(
+            jnp.array(weight),
+            NamedSharding(mesh, P("tensor")),
+        )
     return model
 
 
 def _run_jax(model, input_np, dtype=jnp.float32):
     """Run JAX model and return numpy array."""
-    return np.array(model(jnp.array(input_np, dtype=dtype)))
+    with jax.set_mesh(model.mesh):
+        return np.array(model(jnp.array(input_np, dtype=dtype)))
 
 
 class TestGroupRMSNorm:
@@ -58,18 +87,18 @@ class TestGroupRMSNorm:
     def test_output_shape_matches_input(self):
         """Output shape must match input shape."""
         rng = np.random.default_rng(SEED)
-        input_data = jnp.array(_make_input(rng, (BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE)))
+        input_data = _make_input(rng, (BATCH_SIZE * SEQ_LEN, HIDDEN_SIZE))
 
         model = _make_jax_model()
-        output = model(input_data)
+        output = _run_jax(model, input_data)
 
         assert output.shape == input_data.shape
 
     def test_groups_are_independent(self):
         """Modifying one group must not affect other groups' outputs."""
         rng = np.random.default_rng(SEED)
 
-        input_original = _make_input(rng, (1, 1, HIDDEN_SIZE))
+        input_original = _make_input(rng, (1, HIDDEN_SIZE))
         input_modified = input_original.copy()
         input_modified[..., :GROUP_SIZE] = _make_input(rng, (GROUP_SIZE,))  # perturb group 0 only
 
@@ -93,11 +122,31 @@ def test_groups_are_independent(self):
     def test_weight_participates_in_computation(self):
         """Weight parameter must participate in computation correctly."""
         rng = np.random.default_rng(SEED)
-        input_data = _make_input(rng, (BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE))
+        input_data = _make_input(rng, (BATCH_SIZE * SEQ_LEN, HIDDEN_SIZE))
         weight = _make_weight(rng)
 
         model = _make_jax_model(weight=weight)
         jax_output = _run_jax(model, input_data)
         expected = _numpy_group_rmsnorm_fp64(input_data, weight, NUM_GROUPS, EPSILON)
 
         np.testing.assert_allclose(jax_output, expected, rtol=FP32_RTOL, atol=FP32_ATOL)
+
+    def test_rejects_tp_smaller_than_num_groups(self):
+        """Tensor parallelism must be at least the number of RMS groups."""
+        mesh = Mesh(
+            np.array(jax.devices()[:1]).reshape(1, 1),
+            axis_names=("data", "tensor"),
+            axis_types=(AxisType.Explicit, AxisType.Explicit),
+        )
+
+        with pytest.raises(
+            ValueError,
+            match="tensor parallel size.*num_groups",
+        ):
+            GroupRMSNorm(
+                HIDDEN_SIZE,
+                num_groups=NUM_GROUPS,
+                epsilon=EPSILON,
+                kernel_axes=("tensor",),
+                mesh=mesh,
+            )