diff --git a/gov/work/2026-06-07-rebase-runtime-env-to-torch-2-11-dev-cu13-image-deepgemm-sgl-deep-gemm-port-compat.toml b/gov/work/2026-06-07-rebase-runtime-env-to-torch-2-11-dev-cu13-image-deepgemm-sgl-deep-gemm-port-compat.toml index 221b7bec1..9fa50344e 100644 --- a/gov/work/2026-06-07-rebase-runtime-env-to-torch-2-11-dev-cu13-image-deepgemm-sgl-deep-gemm-port-compat.toml +++ b/gov/work/2026-06-07-rebase-runtime-env-to-torch-2-11-dev-cu13-image-deepgemm-sgl-deep-gemm-port-compat.toml @@ -20,6 +20,31 @@ date = "2026-06-07" scope = "deepgemm-port" content = "Complete-migration sweep (user: 要迁移就完整迁移). Verified EVERY deep_gemm symbol our code calls against HEAD/0.1.2. UNCHANGED+confirmed: fp8_mqa_logits (ks/ke cu_seqlens, not paged ctx - no 2D issue), fp8_gemm_nt, bf16_gemm_nt/nn, get_mk_alignment_for_contiguous_layout (no-arg), transform_sf_into_required_layout (same kwargs), get/set_num_sms, fp8_m_grouped_gemm_nt_masked (HEAD only ADDS optional recipe_a/b NVFP4 + _ensure_cuda, not required for FP8). PORTED beyond the 4 compat edits: (5) #26839 SBO masked-return unpack guard in moe_runner/deep_gemm.py (our line 388 had the unguarded unpack; our launch enables --enable-single-batch-overlap). (6) #23979 PDL-on-by-default in entrypoint.py (hasattr-guarded set_pdl(True), get_bool_env_var SGLANG_DEEPGEMM_PDL default true - perf default matching new wheel). (7) engine.py sglang-kernel assert 0.4.0->0.4.3 (migrated code requires sgl-deep-gemm-era wheel; (N_total,1) ctx would misbehave on old bundled DeepGEMM). SKIPPED with rationale (not on GLM-5.1-FP8/NSA/deepgemm/H200-SM90 path): #26025 fallback-unsupported-shapes (N/A - our tree has no _varlen_deep_gemm_silu_mul_quant/SGLANG_OPT_USE_JIT_EP_ACTIVATION), #25286 Gemma4 triton_scaled_mm scale layout (triton path not deepgemm), #22300 Minimax fp16/flashinfer-trtllm fallback, #26473/#17392 BF16 features, #26496/#24692 SM120/NVFP4, #26238/#25884 dsv4, AMD/ROCm. All 7 edited files py_compile clean. Runtime validation pending the torch-2.11 image (Phase 2/3)." +[[content.journal]] +date = "2026-06-07" +scope = "env-rebase" +content = "Phase-2 prep (tai-kernel rebuild characterization). tai-kernel has NO prebuilt .so / no setup.py — it JIT-compiles its CUDA (.cu/.cc, at::Tensor + TORCH_LIBRARY_FRAGMENT(tai_kernel_ops)) via torch.utils.cpp_extension.load at first import, caching to ~/.cache/tai-kernel/ops/. GOOD: auto-recompiles against the runtime torch (2.11) headers on first import — no manual wheel rebuild. GOTCHA: the cache digest = md5(source bytes + torch.version.cuda + ldflags) does NOT include torch.__version__; torch 2.9.1+cu130 -> 2.11.0+cu130 keeps cuda='13.0' so the digest is unchanged -> a persisted ~/.cache/tai-kernel would be WRONGLY REUSED (ABI-incompatible) against 2.11. Mitigation: fresh container has empty ~/.cache (auto-fresh compile); if cache is on a mounted/persistent volume, rm -rf ~/.cache/tai-kernel before first import. Phase-2 risk = whether the .cu/.cc compile cleanly against torch-2.11 ATen/c10 headers (empirical, verify in image). Phase 1 deep_gemm migration committed 32818784d3." + +[[content.journal]] +date = "2026-06-07" +scope = "env-rebase" +content = "Phase 2/3 VALIDATED on B300 (g0036/g0037, 8x B300 SXM6 each, driver 590.48, image sglang:dev-cu13 = torch 2.11.0+cu130, sgl-kernel 0.4.3, sgl-deep-gemm 0.1.2, flashinfer 0.6.12, transformers 5.8.1, xgrammar 0.2.1, mooncake-transfer-engine-cuda13 0.3.11.post1, deep_ep 1.2.1, sglang-router 0.3.2 - EXACTLY the predicted target). Synced our branch (Phase-1 port) to shared beegfs /mnt/beegfs/syh/sglang-stable (036<->037 shared). In a privileged --gpus all container: pip install -e python + tai-kernel (--no-deps) OK; sglang 0.5.11 editable shadows image default. SMOKE: (a) deep_gemm 0.1.2 wheel has set_pdl=FALSE and get_compile_mode=FALSE -> BOTH my hasattr guards are load-bearing (would crash unguarded); get_paged_mqa/fp8_paged_mqa/contig/masked present. (b) all migrated modules import clean on torch 2.11 (entrypoint w/ PDL edit, nsa_backend w/ _to_2d, fp8_utils, moe_runner/deep_gemm). (c) DEEPGEMM_SCALE_UE8M0=True on Blackwell -> my transform_scale_ue8m0 DLPack stride fix is ACTIVE/load-bearing here (was dormant on H200 SM90). (d) tai-kernel nsa_prefill CUDA ext JIT-compiled vs torch-2.11 ABI in 46s (cleared ~/.cache/tai-kernel first per the digest-gotcha) - the big ABI risk PASSED. Remaining empirical unknowns: model load (transformers 5.8 + GLM-5.1 trust_remote_code), runtime kernels on SM100, mooncake P->D over RoCE roce_rail0..7 (host /sys GIDs read empty; multi-port roce_rail ports1-6 w/ port1 DOWN - mooncake device/port/GID TBD), deepep normal(P)/low_latency(D) on 1-node B300. Topology change: 1-node P (g0036, tp8/ep8) + 1-node D (g0037, dp8); decode drops nnodes2/dist-init. NICs: roce_rail0..7 (GPU rails), bond1 10.25.203.36/37 control plane." + +[[content.journal]] +date = "2026-06-07" +scope = "b300-bringup" +content = "B300 networking diagnosed + RDMA verified. Newly-received cluster; inter-node RDMA appeared dead but was 2 self-inflicted red herrings: (1) per-rail VRFs (CX8_PLB-config.yaml tables 100-107) - main routing table sends 10.33.x out bond1, so naive ping=100% loss; VRF-bound ping (ip vrf exec vrf_vf_rail0 / ping -I eth_vf_rail0) = 0.15ms. (2) container has no iproute2 (ip absent) - irrelevant. Fabric = B300+CX8 dual-plane rail-optimized Spectrum-X RoCE; per-GPU rails roce_vf_rail0..7 RoCE v2 GID idx 3, 10.33.16-23.x/24, 400G, PXB 1:1 GPU (nvidia-smi topo NIC10-17). IB fabric mlx5_12-15 (Quantum-3) also present but rails are the KV path. VERIFIED RDMA 036<->037 in container: ib_write_bw -d roce_vf_rail0 -x 3, handshake to BOND 10.25.203.37, RDMA over rail -> peak 344.93 / AVG 4.53 Gb/s (sustained collapses -> PFC/ECN tuning open issue). Fixed draft scripts (roce_rail->roce_vf_rail0..7). Deliverables on /mnt/beegfs/syh/harness/: run_prefill_b300.sh, run_decode_b300.sh, run_router_b300.sh, b300_rdma_test.sh (host-runnable server/client), b300-bringup-networking.md. Doc also docs_internal/. Memory: reference_b300_cluster. Dispatched 2 opus bg agents: (A) verify prefill per-layer kernel sources on SM100, (B) survey upstream Blackwell/B300/GLM5 perf fixes to consider porting." + +[[content.journal]] +date = "2026-06-07" +scope = "b300-perf" +content = "Two opus agents done. (A) prefill per-layer kernel-source audit on SM100: NO sm_90 lock-in anywhere; all kernels runtime-selected or JIT-built for live arch. Map: NSA attn=sgl_kernel flash_mla_sparse_fwd (heads 64->128 on Blackwell); indexer=deep_gemm fp8_mqa/paged_mqa/metadata (0.1.2); MoE=deep_gemm m_grouped_fp8_gemm_nt_contiguous + TAI three_step; CP shared-KV=tai_kernel CUDA; FP8 KV=Triton; dense block-fp8=deep_gemm fp8_gemm_nt + sgl_kernel ue8m0 quant. Validated the SECOND tai extension (quantization/three_step, the no-torch-fallback MoE hot path) JIT-compiles on cu130/SM100 in 35s (nsa_prefill was 46s). Notable non-correctness: --enable-single-batch-overlap is INERT on this config (Blackwell+deep_gemm+deepep-normal disables masked-down overlap; prefill uses contiguous); BS_GT1_DEBUG != DEBUG_CP_SHARED_KV (TAI fast paths stay active). (B) perf-fix survey -> user decisions: TOOK #25299 (cache MQA-logits budget, drop per-call mem_get_info host sync) + #22232 (data_ptr-guarded copy_ instead of rope clone; q_rope/k_rope are split views of query/key so in-place RoPE -> guard no-op). Verified provenance + correctness first-principles, import-validated on B300, committed 6bc23945e4. SKIPPED per user: #21332 (un-force MHA one-shot on Blackwell - risky), #23856 (torch.mm indexer GEMM). EVALUATE later: #27320 (kv-index parallelize), #26424 (EAGLE topk=1 skip cat/topk/sort), #25859 (MQA mem-ratio knob, pairs w/ #25299)." + +[[content.journal]] +date = "2026-06-07" +scope = "indexcache" +content = "IndexCache (NSA indexer topk reuse across layers, user's 'every-4-layers' feature) = upstream #21405 (5a35316417 'Enable IndexCache for DeepSeek V3.2', 2026-04-05) + #27114 (393d0e169e gate bugfix + index_skip_topk_offset). NOT in our tree (no skip_topk/index_topk_freq plumbing). GLM-5.1-FP8 config.json has index_topk/index_n_heads/index_head_dim but NO index_topk_freq -> defaults to 1 (no sharing); user will set index_topk_freq=4 later (Zhipu confirmed accuracy OK, so no accuracy-validation needed on our side). Mechanism: thread prev_topk_indices through Model.forward loop -> DecoderLayer -> AttentionMLA forward/forward_prepare/forward_absorb_prepare (gate: 'if not skip_topk or (is_nextn and prev is None): indexer() else reuse') -> forward_absorb_core returns (output, topk_indices). Our forward_mla ALREADY threads topk_indices prepare->core via inner_state (350/363), so carrier exists. #24392 (capture/output infra, 18 files) is ORTHOGONAL - HEAD reuse path doesn't depend on it - SKIPPED. Dispatched opus agent in isolated WORKTREE to port (match HEAD form; find+fix ALL DecoderLayer.forward callers for the new 3-tuple return; py_compile). Will review diff + B300-validate before merging to main tree. Agent B perf-survey MISSED #21405 (in range but not surfaced) - survey not exhaustive." + [[content.acceptance_criteria]] text = "Port deep_gemm sgl-deep-gemm 0.1.2 API compat (2D->N_total,1 ctx_lens, compile-mode hasattr guard, positional m_indices, DLPack stride fix, SBO masked unpack guard, PDL default, kernel version assert)" status = "pending" diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index a18c1b765..afbb251eb 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -102,6 +102,7 @@ class DeepseekMLAForwardMixin: forward_batch: ForwardBatch, zero_allocator: BumpAllocator, llama_4_scaling: Optional[torch.Tensor] = None, + prev_topk_indices: Optional[torch.Tensor] = None, ): from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode @@ -193,18 +194,16 @@ class DeepseekMLAForwardMixin: q = self.q_b_proj(q)[0].view( -1, self.num_local_heads, self.qk_head_dim ) - topk_indices = self.indexer( - x=hidden_states, - q_lora=q_lora, - positions=positions, - forward_batch=forward_batch, - layer_id=self.layer_id, - ) - current_stream.wait_stream(self.alt_stream) - else: - k_nope = k_nope.unsqueeze(1) - q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) - if q_lora is not None: + # skip_topk (shared) layers carry no indexer weights in the + # checkpoint, so they must reuse the carried topk and never run + # the indexer. Do NOT widen this to `or prev_topk_indices is + # None` (the original #21405 gate): that recomputes with an + # uninitialized indexer whenever cross-layer propagation is + # unavailable (e.g. the TBO op path drops topk_indices), + # reintroducing the >index_topk garbling. The is_nextn clause is + # the sole intentional fallback (the nextn layer has its own + # weights). + if not self.skip_topk or (self.is_nextn and prev_topk_indices is None): topk_indices = self.indexer( x=hidden_states, q_lora=q_lora, @@ -212,6 +211,28 @@ class DeepseekMLAForwardMixin: forward_batch=forward_batch, layer_id=self.layer_id, ) + else: + topk_indices = prev_topk_indices + current_stream.wait_stream(self.alt_stream) + else: + k_nope = k_nope.unsqueeze(1) + q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + if q_lora is not None: + # See the skip_topk note above: shared layers have no + # indexer weights, so this gate must not fall back to + # computing when prev_topk_indices is None. + if not self.skip_topk or ( + self.is_nextn and prev_topk_indices is None + ): + topk_indices = self.indexer( + x=hidden_states, + q_lora=q_lora, + positions=positions, + forward_batch=forward_batch, + layer_id=self.layer_id, + ) + else: + topk_indices = prev_topk_indices else: q = self.q_proj(hidden_states)[0].view( -1, self.num_local_heads, self.qk_head_dim @@ -547,7 +568,14 @@ class DeepseekMLAForwardMixin: ) output, _ = self.o_proj(attn_bmm_output) - return output + if self.next_skip_topk is None: + return output + + # Return topk_indices for the next layer when enabling index cache + if not self.next_skip_topk: + return output, None + else: + return output, topk_indices def _maybe_write_cp_shared_local_mla_kv( self: DeepseekV2AttentionMLA, diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 66c7db029..715fa6ac2 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -465,7 +465,7 @@ class DeepseekModelNextN(nn.Module): ) residual = None with get_global_expert_distribution_recorder().disable_this_region(): - hidden_states, residual = self.decoder( + hidden_states, residual, topk_indices = self.decoder( positions, hidden_states, forward_batch, diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 1a6e3ef87..720cbe3c3 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1093,9 +1093,11 @@ class DeepseekV2AttentionMLA( prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, skip_rope: bool = False, + is_nextn: bool = False, ) -> None: super().__init__() self.layer_id = layer_id + self.is_nextn = is_nextn self.hidden_size = hidden_size self.qk_nope_head_dim = qk_nope_head_dim self.qk_rope_head_dim = qk_rope_head_dim @@ -1162,6 +1164,8 @@ class DeepseekV2AttentionMLA( prefix=add_prefix("kv_a_proj_with_mqa", prefix), ) + self.skip_topk = None + self.next_skip_topk = None if self.use_nsa: is_neox_style = not getattr(config, "indexer_rope_interleave", False) self.indexer = Indexer( @@ -1182,6 +1186,47 @@ class DeepseekV2AttentionMLA( layer_id=layer_id, alt_stream=alt_stream, ) + # Refer: https://arxiv.org/abs/2603.12201 for more details. + # skip_topk: when True, this layer will skip computation and reuse previous layer's topk indices. + # next_skip_topk: when True, the next layer will skip computation and reuse this layer's topk indices. + if is_nextn: + self.skip_topk = True + self.next_skip_topk = True + else: + self.index_topk_freq = getattr(config, "index_topk_freq", 1) + self.index_topk_pattern = getattr(config, "index_topk_pattern", None) + self.index_skip_topk_offset = getattr( + config, "index_skip_topk_offset", None + ) + if ( + self.index_topk_pattern is None + and self.index_skip_topk_offset is not None + ): + assert self.index_skip_topk_offset > 0, ( + "index_skip_topk_offset must be positive; offset <= 0 " + "marks layer 0 as skip_topk with no prior topk to reuse" + ) + self.skip_topk = ( + max(layer_id - self.index_skip_topk_offset + 1, 0) + % self.index_topk_freq + != 0 + ) + self.next_skip_topk = ( + max(layer_id - self.index_skip_topk_offset + 2, 0) + % self.index_topk_freq + != 0 + ) + elif self.index_topk_pattern is None: + self.skip_topk = max(layer_id - 1, 0) % self.index_topk_freq != 0 + self.next_skip_topk = layer_id % self.index_topk_freq != 0 + else: + self.skip_topk = self.index_topk_pattern[layer_id] == "S" + if layer_id < len(self.index_topk_pattern) - 1: + self.next_skip_topk = ( + self.index_topk_pattern[layer_id + 1] == "S" + ) + else: + self.next_skip_topk = False self.kv_b_proj = ColumnParallelLinear( self.kv_lora_rank, @@ -1313,9 +1358,14 @@ class DeepseekV2AttentionMLA( ) def op_core(self, state): - state.hidden_states_after_attn = self.forward_core( - state.pop("attn_intermediate_state") - ) + result = self.forward_core(state.pop("attn_intermediate_state")) + # forward_core may return (hidden_states, topk_indices) for NSA models + # with index cache enabled. In the TBO path, topk_indices is not + # propagated between layers, so we discard it here. + if isinstance(result, tuple): + state.hidden_states_after_attn = result[0] + else: + state.hidden_states_after_attn = result def forward( self, @@ -1325,6 +1375,7 @@ class DeepseekV2AttentionMLA( zero_allocator: BumpAllocator, layer_scatter_modes: LayerScatterModes = None, llama_4_scaling: Optional[torch.Tensor] = None, + prev_topk_indices: Optional[torch.Tensor] = None, ): s = self.forward_prepare( positions=positions, @@ -1333,6 +1384,7 @@ class DeepseekV2AttentionMLA( zero_allocator=zero_allocator, layer_scatter_modes=layer_scatter_modes, llama_4_scaling=llama_4_scaling, + prev_topk_indices=prev_topk_indices, ) return self.forward_core(s) @@ -1344,6 +1396,7 @@ class DeepseekV2AttentionMLA( zero_allocator: BumpAllocator, layer_scatter_modes: LayerScatterModes = None, llama_4_scaling: Optional[torch.Tensor] = None, + prev_topk_indices: Optional[torch.Tensor] = None, ): if self.attn_mha.kv_b_proj is None: self.attn_mha.kv_b_proj = self.kv_b_proj @@ -1383,7 +1436,12 @@ class DeepseekV2AttentionMLA( ) elif attn_forward_method == AttnForwardMethod.MLA: inner_state = self.forward_absorb_prepare( - positions, hidden_states, forward_batch, zero_allocator, llama_4_scaling + positions, + hidden_states, + forward_batch, + zero_allocator, + llama_4_scaling, + prev_topk_indices, ) elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_ROCM: inner_state = self.forward_absorb_fused_mla_rope_prepare( @@ -1543,6 +1601,7 @@ class DeepseekV2DecoderLayer(nn.Module): reduce_results=False, prefix=add_prefix("self_attn", prefix), alt_stream=alt_stream, + is_nextn=is_nextn, ) if not hasattr(config, "q_lora_rank") and envs.SGLANG_USE_AG_AFTER_QLORA.get(): raise ValueError( @@ -1645,6 +1704,7 @@ class DeepseekV2DecoderLayer(nn.Module): zero_allocator: BumpAllocator, gemm_output_zero_allocator: BumpAllocator = None, llama_4_scaling: Optional[torch.Tensor] = None, + prev_topk_indices: Optional[torch.Tensor] = None, ) -> torch.Tensor: quant_format = ( "mxfp4" @@ -1694,6 +1754,7 @@ class DeepseekV2DecoderLayer(nn.Module): zero_allocator=zero_allocator, llama_4_scaling=llama_4_scaling, layer_scatter_modes=self.layer_scatter_modes, + prev_topk_indices=prev_topk_indices, ) finally: if previous_cp_shared_kv_num_model_layers is None: @@ -1703,6 +1764,13 @@ class DeepseekV2DecoderLayer(nn.Module): previous_cp_shared_kv_num_model_layers ) + # forward returns (hidden_states, topk_indices) for NSA models with the + # index cache enabled; otherwise a plain hidden_states tensor. + if isinstance(hidden_states, tuple): + hidden_states, topk_indices = hidden_states + else: + topk_indices = None + hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states, residual, forward_batch ) @@ -1738,7 +1806,7 @@ class DeepseekV2DecoderLayer(nn.Module): ) self._notify_cp_hicache_layer_end(forward_batch) - return hidden_states, residual + return hidden_states, residual, topk_indices def op_comm_prepare_attn( self, @@ -2018,6 +2086,7 @@ class DeepseekV2Model(nn.Module): elif self.first_k_dense_replace < normal_start_layer: normal_end_layer = normal_start_layer = 0 aux_hidden_states = [] + topk_indices = None for i in range(normal_start_layer, normal_end_layer): # NOTE: torch dynamo does not support graph break in context manager ctx = ( @@ -2035,7 +2104,7 @@ class DeepseekV2Model(nn.Module): else: aux_hidden_states.append(hidden_states + residual) layer = self.layers[i] - hidden_states, residual = layer( + hidden_states, residual, topk_indices = layer( positions, hidden_states, forward_batch, @@ -2043,6 +2112,7 @@ class DeepseekV2Model(nn.Module): zero_allocator, gemm_output_zero_allocator, llama_4_scaling, + prev_topk_indices=topk_indices, ) if normal_end_layer != self.end_layer: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 67eaad907..ddfa97df6 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1584,6 +1584,25 @@ class ServerArgs: self.attention_backend = "nsa" logger.info("Use nsa attention backend for DeepSeek with DSA.") + # IndexCache (NSA/DSA indexer topk sharing) reuses one layer's topk + # across the following layers, threaded through the monolithic decoder + # forward. The two-batch-overlap op path does NOT propagate topk_indices + # between layers, so shared (skip_topk) layers would run sparse attention + # with no indices. Reject the combination. (upstream PR #27114) + index_topk_freq = getattr(hf_config, "index_topk_freq", 1) + index_topk_pattern = getattr(hf_config, "index_topk_pattern", None) + if self.enable_two_batch_overlap and ( + index_topk_freq > 1 + or (index_topk_pattern is not None and "S" in index_topk_pattern) + ): + raise ValueError( + "--enable-two-batch-overlap is not supported with NSA/DSA " + "index-topk sharing (index_topk_freq > 1 or an " + "index_topk_pattern containing shared layers): the TBO op " + "path does not propagate topk indices across layers, so " + "shared layers would run sparse attention without indices." + ) + if not is_npu(): # CUDA or ROCm GPU if self.enable_nsa_prefill_context_parallel: logger.warning(