[qwen 3.8 next] Fuse Mamba tracking index lookup and skip uniform-position transfers

python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py

@@ -893,3 +893,58 @@ def track_mamba_states_all_layers(
         BLOCK_SIZE,
         check_freed_slots,
     )
+
+
+@triton.jit
+def _gather_mamba_track_indices_kernel(
+    mapping,
+    requests,
+    positions,
+    output,
+    N: tl.constexpr,
+    ROWS: tl.constexpr,
+    ROW_STRIDE: tl.constexpr,
+    COL_STRIDE: tl.constexpr,
+    REQ_STRIDE: tl.constexpr,
+    POS_STRIDE: tl.constexpr,
+    BLOCK: tl.constexpr,
+    HAS_POSITIONS: tl.constexpr,
+    UNIFORM_POSITION,
+):
+    i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    req = tl.load(requests + i * REQ_STRIDE, i < N, other=0).to(tl.int64)
+    req = tl.where(req < 0, req + ROWS, req)
+    if HAS_POSITIONS:
+        pos = tl.load(positions + i * POS_STRIDE, i < N, other=0).to(tl.int64)
+    else:
+        pos = tl.full((), 0, tl.int64) + UNIFORM_POSITION
+    value = tl.load(mapping + req * ROW_STRIDE + pos * COL_STRIDE, i < N, other=0)
+    tl.store(output + i, value, i < N)
+
+
+def gather_mamba_track_indices(
+    mapping, requests, positions=None, uniform_position=None
+):
+    n = requests.numel()
+    if positions is not None:
+        assert positions.numel() == n
+    else:
+        assert uniform_position is not None and 0 <= uniform_position < mapping.shape[1]
+    output = torch.empty((n,), dtype=torch.int64, device=mapping.device)
+    if n:
+        _gather_mamba_track_indices_kernel[(triton.cdiv(n, 128),)](
+            mapping,
+            requests,
+            positions,
+            output,
+            n,
+            mapping.shape[0],
+            mapping.stride(0),
+            mapping.stride(1),
+            requests.stride(0),
+            positions.stride(0) if positions is not None else 0,
+            BLOCK=128,
+            HAS_POSITIONS=positions is not None,
+            UNIFORM_POSITION=uniform_position if positions is None else 0,
+        )
+    return output

python/sglang/srt/managers/schedule_batch.py

@@ -2188,9 +2188,7 @@ def set_mamba_track_indices_from_reqs(
     spec track plan, see mamba_lazy_spec_prepare).
     """
     req_to_token_pool = batch.req_to_token_pool
-    all_buffers = req_to_token_pool.req_index_to_mamba_ping_pong_track_buffer_mapping[
-        batch.req_pool_indices
-    ]  # (bs, ping_pong_size), int64, on device
+    mapping = req_to_token_pool.req_index_to_mamba_ping_pong_track_buffer_mapping
     if track_positions is None:
         # Guard: mamba_next_track_idx may be None for requests that haven't
         # gone through _alloc_ping_pong_buffer yet (e.g., spec v2 verify path).
@@ -2204,18 +2202,42 @@ def set_mamba_track_indices_from_reqs(
             for req in batch.reqs
         ]
     batch.mamba_track_buffer_indices = list(track_positions)
+    if (
+        mapping.is_cuda
+        and track_positions
+        and all(pos == track_positions[0] for pos in track_positions)
+    ):
+        from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
+            gather_mamba_track_indices,
+        )
+
+        batch.mamba_track_indices = gather_mamba_track_indices(
+            mapping, batch.req_pool_indices, uniform_position=track_positions[0]
+        )
+        return
+
     idx = (
         torch.tensor(
             track_positions,
             dtype=torch.int64,
             pin_memory=True,
         )
         .unsqueeze(1)
-        .to(device=all_buffers.device, non_blocking=True)
-    )
-    batch.mamba_track_indices = (
-        torch.gather(all_buffers, 1, idx).squeeze(1).to(torch.int64)
+        .to(device=mapping.device, non_blocking=True)
     )
+    if mapping.is_cuda:
+        from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
+            gather_mamba_track_indices,
+        )
+
+        batch.mamba_track_indices = gather_mamba_track_indices(
+            mapping, batch.req_pool_indices, idx.squeeze(1)
+        )
+    else:
+        all_buffers = mapping[batch.req_pool_indices]
+        batch.mamba_track_indices = (
+            torch.gather(all_buffers, 1, idx).squeeze(1).to(torch.int64)
+        )
 
 
 def release_req(

test/registered/unit/layers/test_next_tracking.py

@@ -0,0 +1,85 @@
+import pytest
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
+
+
+@pytest.mark.parametrize("bs", [0, 1, 3, 129])
+@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
+def test_mamba_track_gather_graph(bs, dtype):
+    from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
+        gather_mamba_track_indices,
+    )
+
+    mapping = torch.arange(300 * 12, device="cuda", dtype=dtype).reshape(300, 12)[
+        ::2, ::3
+    ]
+    requests = torch.arange(bs * 2, device="cuda", dtype=dtype)[::2] % 150
+    requests -= 150
+    positions = torch.arange(bs * 2, device="cuda", dtype=dtype)[::2] % 4
+    gather_mamba_track_indices(mapping, requests, positions)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        actual = gather_mamba_track_indices(mapping, requests, positions)
+    for i in range(3):
+        positions.add_(1).remainder_(4)
+        mapping.add_(i)
+        graph.replay()
+        expected = (
+            mapping[requests.long()]
+            .gather(1, positions.long().unsqueeze(1))
+            .squeeze(1)
+            .long()
+        )
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("bs", [0, 1, 3, 129])
+@pytest.mark.parametrize("position", [0, 1, 3])
+def test_mamba_uniform_track_gather_graph(bs, position):
+    from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
+        gather_mamba_track_indices,
+    )
+
+    mapping = torch.arange(300 * 12, device="cuda").reshape(300, 12)[::2, ::3]
+    requests = torch.arange(bs * 2, device="cuda")[::2] % 150
+    gather_mamba_track_indices(mapping, requests, uniform_position=position)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        actual = gather_mamba_track_indices(
+            mapping, requests, uniform_position=position
+        )
+    for i in range(3):
+        requests.add_(1).remainder_(150)
+        mapping.add_(i)
+        graph.replay()
+        torch.testing.assert_close(actual, mapping[requests, position], rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("positions", [[0], [1], [0, 0, 0], [0, 1, 0]])
+@pytest.mark.parametrize("override", [False, True])
+def test_mamba_track_indices_from_reqs(positions, override):
+    from types import SimpleNamespace
+
+    from sglang.srt.managers.schedule_batch import set_mamba_track_indices_from_reqs
+
+    mapping = torch.arange(20, device="cuda", dtype=torch.int64).reshape(10, 2)
+    requests = torch.arange(len(positions), device="cuda", dtype=torch.int64) + 2
+    batch = SimpleNamespace(
+        req_to_token_pool=SimpleNamespace(
+            req_index_to_mamba_ping_pong_track_buffer_mapping=mapping
+        ),
+        req_pool_indices=requests,
+        reqs=[
+            SimpleNamespace(kv=SimpleNamespace(mamba_next_track_idx=(p if p else None)))
+            for p in positions
+        ],
+    )
+    set_mamba_track_indices_from_reqs(
+        batch, track_positions=positions if override else None
+    )
+    expected = mapping[requests, torch.tensor(positions, device="cuda")]
+    torch.testing.assert_close(batch.mamba_track_indices, expected, rtol=0, atol=0)
+    assert batch.mamba_track_buffer_indices == positions