Fix lint error (sgl-project#427)

Author: silencejade

python/sgl_kernel_npu/sgl_kernel_npu/mamba/mamba_state_update_triton.py

@@ -113,8 +113,12 @@ def move_intermediate_cache(
     dst_layer_stride, dst_size_stride = int(ssm_states.stride()[0]), int(
         ssm_states.stride()[1]
     )
-    assert len(dst_indices_tensor) == len(last_steps_tensor), "Destination indices lengths must match"
-    assert len(src_indices_tensor) == len(last_steps_tensor), "Source indices lengths must match"
+    assert len(dst_indices_tensor) == len(
+        last_steps_tensor
+    ), "Destination indices lengths must match"
+    assert len(src_indices_tensor) == len(
+        last_steps_tensor
+    ), "Source indices lengths must match"
 
     # Grid: one thread per valid index
     grid = (len(dst_indices_tensor),)

tests/python/sgl_kernel_npu/test_mamba_state_update.py

@@ -139,7 +139,9 @@ def test_move_intermediate_cache(
     valid_indices = random.sample(population, num_valid)
     last_step_pos = [random.randint(0, D - 1) for _ in range(num_valid)]
     dst_indices_tensor = torch.tensor(valid_indices, device=device, dtype=torch.int32)
-    src_indices_tensor = torch.arange(dst_indices_tensor.shape[0], device=device, dtype=torch.int32)
+    src_indices_tensor = torch.arange(
+        dst_indices_tensor.shape[0], device=device, dtype=torch.int32
+    )
     last_steps_tensor = torch.tensor(last_step_pos, device=device, dtype=torch.int32)
 
     valid_mask = last_steps_tensor >= 0
@@ -151,6 +153,12 @@ def test_move_intermediate_cache(
         :, src_state_indices, valid_last_steps
     ]
 
-    move_intermediate_cache(dst_cache_clone, src_cache, dst_indices_tensor, src_indices_tensor, last_steps_tensor)
+    move_intermediate_cache(
+        dst_cache_clone,
+        src_cache,
+        dst_indices_tensor,
+        src_indices_tensor,
+        last_steps_tensor,
+    )
 
     assert_close("move_cache", dst_cache, dst_cache_clone, 1e-3)