From 050dbaca120f90171b141b52d3ff26c6fbcd66a6 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Mon, 13 Jul 2026 02:38:34 +0000 Subject: [PATCH 01/25] init --- .../deepseek_v4_backend_hip_radix.py | 119 ++++++++++++++---- python/sglang/srt/speculative/dflash_utils.py | 29 ++++- 2 files changed, 121 insertions(+), 27 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 1755e8212fe1..2c413702d497 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -414,6 +414,7 @@ class DeepseekV4HipRadixBackend( # both children and leaks ROCm HSA resources (HSA_STATUS_ERROR_OUT_OF_RESOURCES). # TboAttnBackend reads this to skip children in the *_graph paths only. tbo_supports_cuda_graph = False + supports_ragged_verify_graph: bool = True def __init__( self, @@ -458,6 +459,12 @@ def __init__( self.speculative_num_draft_tokens: int = ( model_runner.server_args.speculative_num_draft_tokens ) + if ( + self.speculative_num_draft_tokens is not None + and getattr(model_runner, "is_draft_worker", False) + and model_runner.spec_algorithm.is_dspark() + ): + self.speculative_num_draft_tokens -= 1 self.speculative_step_id = speculative_step_id self.forward_metadata: Union[ DSV4Metadata, @@ -534,14 +541,34 @@ def init_forward_metadata_prefill( extend_seq_lens_cpu: List[int], need_compress: bool = True, use_prefill_cuda_graph: bool = False, + compress_gpu_plan: bool = False, + extend_start_loc: Optional[torch.Tensor] = None, ) -> DSV4Metadata: - seq_lens_casual, req_pool_indices_repeated = self.expand_prefill_casually( - num_tokens=num_tokens, - seq_lens=seq_lens_cpu, - extend_seq_lens=extend_seq_lens_cpu, - req_pool_indices=req_pool_indices, - padded_num_tokens=out_cache_loc.shape[0], - ) + if extend_start_loc is not None: + from sglang.srt.layers.attention.dsv4.attn_metadata_kernels import ( + ExpandPrefillCausally, + ) + + _expanded = ExpandPrefillCausally.execute( + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + extend_start_loc=extend_start_loc, + seq_lens_cpu=None, + extend_seq_lens_cpu=None, + num_tokens=num_tokens, + padded_num_tokens=out_cache_loc.shape[0], + ) + seq_lens_casual = _expanded.seq_lens_casual + req_pool_indices_repeated = _expanded.req_pool_indices_repeated + else: + seq_lens_casual, req_pool_indices_repeated = self.expand_prefill_casually( + num_tokens=num_tokens, + seq_lens=seq_lens_cpu, + extend_seq_lens=extend_seq_lens_cpu, + req_pool_indices=req_pool_indices, + padded_num_tokens=out_cache_loc.shape[0], + ) core_attn_metadata = self.make_core_attn_metadata( req_to_token=self.req_to_token, req_pool_indices_repeated=req_pool_indices_repeated, @@ -561,6 +588,20 @@ def init_forward_metadata_prefill( ) if not need_compress: create = _create_dummy_paged_compress_data + elif compress_gpu_plan: + create = functools.partial( + create_paged_compressor_data, + is_prefill=True, + token_to_kv_pool=self.token_to_kv_pool, + req_to_token=self.req_to_token, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + seq_lens_cpu=None, + extend_lens=extend_seq_lens, + extend_lens_cpu=None, + num_q_tokens=num_tokens, + use_prefill_cuda_graph=use_prefill_cuda_graph, + ) else: create = functools.partial( create_paged_compressor_data, @@ -590,6 +631,7 @@ def init_forward_metadata_target_verify( extend_seq_lens: Optional[torch.Tensor] = None, use_prefill_cuda_graph: bool = False, seq_lens_cpu: Optional[List[int]] = None, + ragged_layout=None, ) -> Union[DSV4Metadata, DSV4RawVerifyMetadata]: # HIP path: build target-verify metadata eagerly even when # SGLANG_PREP_IN_CUDA_GRAPH is enabled. The raw/lazy-upgrade route can @@ -603,6 +645,7 @@ def init_forward_metadata_target_verify( seq_lens_cpu=seq_lens_cpu, out_cache_loc=out_cache_loc, use_prefill_cuda_graph=use_prefill_cuda_graph, + ragged_layout=ragged_layout, ) def init_forward_metadata_target_verify_old( @@ -613,13 +656,36 @@ def init_forward_metadata_target_verify_old( seq_lens_cpu: Optional[List[int]] = None, out_cache_loc: Optional[torch.Tensor] = None, use_prefill_cuda_graph: bool = False, + ragged_layout=None, ) -> DSV4Metadata: batch_size = len(seq_lens) - seq_lens = seq_lens + self.speculative_num_draft_tokens - seq_lens_cpu = [x + self.speculative_num_draft_tokens for x in seq_lens_cpu] - extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * batch_size - extend_seq_lens = self._move_to_device(extend_seq_lens_cpu) - num_tokens = self.speculative_num_draft_tokens * batch_size + extend_start_loc = None + if ragged_layout is not None: + verify_lens_dev = ragged_layout.verify_lens.to( + device=seq_lens.device, dtype=torch.int32 + ) + extend_start_loc = ragged_layout.extend_start_loc.to( + device=seq_lens.device, dtype=torch.int32 + ) + extend_seq_lens = verify_lens_dev + seq_lens = seq_lens + verify_lens_dev.to(seq_lens.dtype) + # Total verify tokens to expand. For the graph path the padded layout + # sets total_verify_tokens == graph_num_tokens (tier); the eager path + # resolves a device-assembled layout whose total_verify_tokens is None, + # so fall back to sum(verify_lens) (== real total; padded == tier). + num_tokens = ragged_layout.total_verify_tokens + if num_tokens is None: + num_tokens = int(verify_lens_dev.sum().item()) + else: + num_tokens = int(num_tokens) + extend_seq_lens_cpu = None + seq_lens_cpu = None + else: + seq_lens = seq_lens + self.speculative_num_draft_tokens + seq_lens_cpu = [x + self.speculative_num_draft_tokens for x in seq_lens_cpu] + extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * batch_size + num_tokens = self.speculative_num_draft_tokens * batch_size + extend_seq_lens = self._move_to_device(extend_seq_lens_cpu) if out_cache_loc is None: out_cache_loc = seq_lens.new_zeros(num_tokens) return self.init_forward_metadata_prefill( @@ -633,6 +699,8 @@ def init_forward_metadata_target_verify_old( extend_seq_lens_cpu=extend_seq_lens_cpu, need_compress=True, use_prefill_cuda_graph=use_prefill_cuda_graph, + compress_gpu_plan=ragged_layout is not None, + extend_start_loc=extend_start_loc, ) def make_forward_metadata_from_raw_verify( @@ -848,6 +916,8 @@ def init_forward_metadata_out_graph( chosen_max_seq_len = self.MAX_SEQ_LEN_FOR_CAPTURE assert actual_max_seq_len <= chosen_max_seq_len + graph_key = bs + if bucket == _GraphBucket.DECODE_OR_IDLE: assert out_cache_loc is not None assert len(out_cache_loc.shape) == 1, f"{out_cache_loc.shape=}" @@ -864,14 +934,14 @@ def init_forward_metadata_out_graph( out_cache_loc=out_cache_loc_padded, ) elif bucket == _GraphBucket.TARGET_VERIFY: - if resolve_ragged_verify_layout(forward_batch) is not None: - raise NotImplementedError( - "DSV4 ragged verify is not supported on the HIP backend " - "(DeepseekV4HipRadixBackend) cuda-graph path; disable " - "SGLANG_RAGGED_VERIFY_MODE or use a CUDA device." - ) assert out_cache_loc is not None - num_tokens_v = self.speculative_num_draft_tokens * bs + ragged_layout = resolve_ragged_verify_layout(forward_batch) + if ragged_layout is not None: + ragged_layout = ragged_layout.padded_to_bucket(padded_bs=bs) + num_tokens_v = ragged_layout.graph_num_tokens + graph_key = num_tokens_v + else: + num_tokens_v = self.speculative_num_draft_tokens * bs out_cache_loc_padded = torch.nn.functional.pad( out_cache_loc, pad=(0, num_tokens_v - len(out_cache_loc)), @@ -887,6 +957,7 @@ def init_forward_metadata_out_graph( # CPU mirror already available here (== seq_lens, no D2H); # pass it so target_verify skips the per-iter seq_lens.tolist() sync. seq_lens_cpu=seq_lens_cpu.tolist(), + ragged_layout=ragged_layout, ) elif bucket == _GraphBucket.DRAFT_EXTEND: num_tokens_per_bs = self.draft_extend_num_tokens_per_bs @@ -912,7 +983,7 @@ def init_forward_metadata_out_graph( raise NotImplementedError self.replay_cuda_graph_metadata_from( - bs=bs, temp_metadata=temp_metadata, bucket=bucket + bs=graph_key, temp_metadata=temp_metadata, bucket=bucket ) if in_capture: @@ -957,12 +1028,7 @@ def init_forward_metadata(self, forward_batch: ForwardBatch) -> None: out_cache_loc=out_cache_loc, ) elif forward_batch.forward_mode.is_target_verify(): - if resolve_ragged_verify_layout(forward_batch) is not None: - raise NotImplementedError( - "DSV4 ragged verify is not supported on the HIP backend " - "(DeepseekV4HipRadixBackend); disable SGLANG_RAGGED_VERIFY_MODE " - "or use a CUDA device." - ) + ragged_layout = resolve_ragged_verify_layout(forward_batch) metadata = self.init_forward_metadata_target_verify( max_seq_len=max_seq_len, req_pool_indices=req_pool_indices, @@ -972,6 +1038,7 @@ def init_forward_metadata(self, forward_batch: ForwardBatch) -> None: seq_lens_cpu=( seq_lens_cpu.tolist() if seq_lens_cpu is not None else None ), + ragged_layout=ragged_layout, ) elif forward_batch.forward_mode.is_prefill(include_draft_extend_v2=True): extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 11d3938ae23f..c554162bf2a7 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -12,7 +12,7 @@ from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod from sglang.srt.layers.sampler import apply_custom_logit_processor from sglang.srt.managers.schedule_batch import Req -from sglang.srt.utils import is_cuda, is_musa +from sglang.srt.utils import is_cuda, is_hip, is_musa DEFAULT_DFLASH_MASK_TOKEN = "<|MASK|>" @@ -45,6 +45,33 @@ top_k_renorm_prob = None top_p_renorm_prob = None tree_speculative_sampling_target_only = None +elif is_hip(): + + def _dspark_top_k_renorm_prob_torch(probs, top_ks): + top_ks = top_ks.to(device=probs.device) + sorted_probs, sorted_indices = torch.sort(probs, dim=-1, descending=True) + rank = torch.arange(probs.shape[-1], device=probs.device).view(1, -1) + sorted_probs = torch.where( + rank < top_ks.view(-1, 1), sorted_probs, torch.zeros_like(sorted_probs) + ) + renorm = torch.zeros_like(probs).scatter_(-1, sorted_indices, sorted_probs) + return renorm / renorm.sum(dim=-1, keepdim=True).clamp_min(1e-12) + + def _dspark_top_p_renorm_prob_torch(probs, top_ps): + top_ps = top_ps.to(device=probs.device) + sorted_probs, sorted_indices = torch.sort(probs, dim=-1, descending=True) + cumsum_excl = torch.cumsum(sorted_probs, dim=-1) - sorted_probs + sorted_probs = torch.where( + cumsum_excl <= top_ps.view(-1, 1), + sorted_probs, + torch.zeros_like(sorted_probs), + ) + renorm = torch.zeros_like(probs).scatter_(-1, sorted_indices, sorted_probs) + return renorm / renorm.sum(dim=-1, keepdim=True).clamp_min(1e-12) + + top_k_renorm_prob = _dspark_top_k_renorm_prob_torch + top_p_renorm_prob = _dspark_top_p_renorm_prob_torch + tree_speculative_sampling_target_only = None else: top_k_renorm_prob = None top_p_renorm_prob = None From d3de45de68ebf163d66b2300a659115b3a67e2b1 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Tue, 14 Jul 2026 00:31:45 +0000 Subject: [PATCH 02/25] support unified kv --- .../dsv4/unified_kv_kernels/runtime.py | 50 +++++++++ .../srt/mem_cache/deepseek_v4_memory_pool.py | 28 +++++ .../sglang/srt/models/deepseek_v4_dspark.py | 25 ++++- .../dspark_components/dspark_kv_inject.py | 106 +++++++++++++++--- .../dspark_components/dspark_verify.py | 38 +++++-- .../dspark_components/dspark_worker_v2.py | 19 ++++ .../kernels/dspark_verify_window.py | 31 +++++ 7 files changed, 272 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/runtime.py b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/runtime.py index 437fc3e8bfd2..2c6a5677a963 100644 --- a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/runtime.py +++ b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/runtime.py @@ -110,6 +110,56 @@ def store_swa_into_unified( ) +@triton.jit +def _scatter_loc_kernel( + kv_ptr, # [T, D] bf16 + loc_ptr, # [T] int (unified row index; <0 => skip) + unified_ptr, # [pages, D] bf16 + n_rows, + D: tl.constexpr, + BLOCK_D: tl.constexpr, +): + row = tl.program_id(0) + if row >= n_rows: + return + loc = tl.load(loc_ptr + row).to(tl.int64) + if loc < 0: + return + offs = tl.arange(0, BLOCK_D) + mask = offs < D + vals = tl.load(kv_ptr + row * D + offs, mask=mask, other=0.0) + tl.store(unified_ptr + loc * D + offs, vals, mask=mask) + + +def scatter_bf16_into_unified( + *, + kv: torch.Tensor, # [T, head_dim] bf16 (already norm+rope'd) + loc: torch.Tensor, # [T] int32/int64 unified ring row; <0 => skip + unified_kv: torch.Tensor, # [pages, head_dim] bf16 +) -> None: + """Scatter already-norm+rope'd bf16 K into ``unified_kv[loc]`` (skip loc < 0). + + Companion to ``store_swa_into_unified`` for callers that already hold the + precomputed ring row index (the DSpark draft: ``get_unified_swa_loc`` for the + draft forward, or the commit-inject layout for target-hidden injection) and + need per-row commit masking expressed as ``loc == -1``. + """ + n_rows, D = kv.shape + if n_rows == 0: + return + assert kv.is_contiguous() and kv.dtype == unified_kv.dtype + assert loc.is_contiguous() + _scatter_loc_kernel[(n_rows,)]( + kv, + loc, + unified_kv, + n_rows, + D=D, + BLOCK_D=triton.next_power_of_2(D), + num_warps=8, + ) + + # --------------------------------------------------------------------------- # Ragged indptr helper (shared by the decode streams + prefill builders) # --------------------------------------------------------------------------- diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 71ad833b1e04..c75438c8010d 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -1195,6 +1195,34 @@ def set_swa_key_buffer_radix_fused_norm_rope( page_size=self.swa_kv_pool.page_size, ) + def set_unified_key_buffer_radix_fused_norm_rope( + self, + layer_id: int, + swa_loc: torch.Tensor, + kv: torch.Tensor, + kv_weight: torch.Tensor, + eps: float, + freqs_cis: torch.Tensor, + positions: torch.Tensor, + ) -> None: + """unified_kv counterpart of set_swa_key_buffer_radix_fused_norm_rope. + + Under unified_kv the (fp8, paged) swa_kv_pool is None -- SWA K lives in + the shared bf16 unified_kv ring instead. Norm+RoPE the draft KV in place + (the same freqs_cis path the main model uses via _compute_kv_bf16) and + scatter it into ``unified_kv[swa_loc]``. Rows with swa_loc < 0 + (uncommitted verify tokens) are skipped by the scatter. + """ + from sglang.jit_kernel.dsv4 import fused_norm_rope_inplace + from sglang.srt.layers.attention.dsv4.unified_kv_kernels import runtime + + fused_norm_rope_inplace(kv, kv_weight, eps, freqs_cis, positions) + runtime.scatter_bf16_into_unified( + kv=kv, + loc=swa_loc, + unified_kv=self.get_unified_kv(layer_id), + ) + def set_extra_key_buffer_fused( self, layer_id: int, diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index 98a411afd18d..9d52b9cb774a 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -141,6 +141,20 @@ def _store_block_kv( attn_backend, pool: DeepSeekV4TokenToKVPool, ) -> None: + if pool._unified_kv: + # unified_kv: SWA K lives in the shared bf16 ring (swa_kv_pool is + # None). Use the unified ring write target -- get_unified_swa_loc + # recomputes it from live positions for multi-step draft decode. + pool.set_unified_key_buffer_radix_fused_norm_rope( + layer_id=self.layer_id, + swa_loc=attn_backend.get_unified_swa_loc(forward_batch), + kv=kv, + kv_weight=self.kv_norm.weight.data, + eps=self.eps, + freqs_cis=self.freqs_cis, + positions=positions, + ) + return pool.set_swa_key_buffer_radix_fused_norm_rope( layer_id=self.layer_id, swa_loc=attn_backend.get_swa_out_cache_loc(forward_batch), @@ -665,9 +679,18 @@ def write_target_hidden_kv( main_x=main_x, wkv_linears=[stage.self_attn.wkv for stage in self.stages], ) + # Under unified_kv the swa_kv_pool is None; the caller passes a unified + # ring loc (state_slot * ring + pos % ring, -1 for uncommitted) so the + # store just needs to target the bf16 ring instead of the fp8 flashmla + # buffer. Same swa_loc/positions contract either way. + store_kv = ( + pool.set_unified_key_buffer_radix_fused_norm_rope + if pool._unified_kv + else pool.set_swa_key_buffer_radix_fused_norm_rope + ) for stage, kv in zip(self.stages, kvs): attn = stage.self_attn - pool.set_swa_key_buffer_radix_fused_norm_rope( + store_kv( layer_id=attn.layer_id, swa_loc=swa_loc, kv=kv, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py index 329c3f455f35..4855eb633bbe 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py @@ -6,6 +6,7 @@ from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.speculative.dspark_components.kernels.dspark_verify_window import ( BuildCommitInjectLayout, + build_unified_commit_inject_layout, ) from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout @@ -36,6 +37,8 @@ def inject_target_hidden( positions: torch.Tensor, cache_loc_2d: Optional[torch.Tensor] = None, commit_lens: Optional[torch.Tensor] = None, + state_slot: Optional[torch.Tensor] = None, + final_pos: Optional[torch.Tensor] = None, ) -> None: if target_hidden is None or target_hidden.numel() == 0: return @@ -54,6 +57,14 @@ def inject_target_hidden( commit_lens = commit_lens.to( device=device, dtype=torch.int32, non_blocking=True ) + if state_slot is not None: + state_slot = state_slot.to( + device=device, dtype=torch.int64, non_blocking=True + ) + if final_pos is not None: + final_pos = final_pos.to( + device=device, dtype=torch.int64, non_blocking=True + ) pool = self.draft_model_runner.token_to_kv_pool if hasattr(pool, "set_swa_key_buffer_radix_fused_norm_rope"): @@ -64,6 +75,8 @@ def inject_target_hidden( positions=positions, cache_loc_2d=cache_loc_2d, commit_lens=commit_lens, + state_slot=state_slot, + final_pos=final_pos, ) return @@ -86,13 +99,29 @@ def _inject_mla( positions: torch.Tensor, cache_loc_2d: Optional[torch.Tensor], commit_lens: Optional[torch.Tensor], + state_slot: Optional[torch.Tensor] = None, + final_pos: Optional[torch.Tensor] = None, ) -> None: - swa_loc = pool.translate_loc_from_full_to_swa(cache_loc).to(torch.int32) - if commit_lens is not None and cache_loc_2d is not None: - bs, verify_len = cache_loc_2d.shape - col = torch.arange(verify_len, device=cache_loc.device).view(1, -1) - committed_mask = (col < commit_lens.to(torch.long).view(-1, 1)).reshape(-1) - swa_loc = torch.where(committed_mask, swa_loc, torch.full_like(swa_loc, -1)) + if getattr(pool, "_unified_kv", False): + swa_loc = self._unified_inject_loc( + pool=pool, + positions=positions, + cache_loc_2d=cache_loc_2d, + commit_lens=commit_lens, + state_slot=state_slot, + final_pos=final_pos, + ) + else: + swa_loc = pool.translate_loc_from_full_to_swa(cache_loc).to(torch.int32) + if commit_lens is not None and cache_loc_2d is not None: + bs, verify_len = cache_loc_2d.shape + col = torch.arange(verify_len, device=cache_loc.device).view(1, -1) + committed_mask = (col < commit_lens.to(torch.long).view(-1, 1)).reshape( + -1 + ) + swa_loc = torch.where( + committed_mask, swa_loc, torch.full_like(swa_loc, -1) + ) with torch.inference_mode(): self.draft_model.write_target_hidden_kv( @@ -102,6 +131,43 @@ def _inject_mla( pool=pool, ) + def _unified_inject_loc( + self, + *, + pool, + positions: torch.Tensor, + cache_loc_2d: Optional[torch.Tensor], + commit_lens: Optional[torch.Tensor], + state_slot: Optional[torch.Tensor], + final_pos: Optional[torch.Tensor], + ) -> torch.Tensor: + """Ring row for target-hidden injection under unified_kv. + + loc = state_slot * ring + pos % ring, with two skip (-1) rules: + * SWA window: only the last ``win`` tokens per req land in the ring; + older tokens share a ring slot (pos % ring) and would race, so drop + them (needed for long prefill chunks). + * commit gate: uncommitted verify tokens (col >= commit_len) are dropped. + """ + if state_slot is None: + raise RuntimeError( + "unified_kv target-hidden injection requires state_slot " + "(per-token draft req_pool_indices)." + ) + ring = pool.unified_swa_ring_size + win = pool.unified_swa_window + pos = positions.to(torch.int64) + loc = state_slot.to(torch.int64) * ring + pos % ring + if final_pos is not None: + keep = pos > (final_pos.to(torch.int64) - win) + loc = torch.where(keep, loc, torch.full_like(loc, -1)) + if commit_lens is not None and cache_loc_2d is not None: + bs, verify_len = cache_loc_2d.shape + col = torch.arange(verify_len, device=positions.device).view(1, -1) + committed = (col < commit_lens.to(torch.long).view(-1, 1)).reshape(-1) + loc = torch.where(committed, loc, torch.full_like(loc, -1)) + return loc.to(torch.int32) + def inject_ragged( self, *, @@ -119,15 +185,25 @@ def inject_ragged( if hasattr(pool, "set_swa_key_buffer_radix_fused_norm_rope"): if hidden_strided.numel() == 0: return - inject_layout = BuildCommitInjectLayout.execute( - req_pool_indices=batch.req_pool_indices, - req_to_token=self.model_runner.req_to_token_pool.req_to_token, - prefix_lens=prefix_lens, - block_pos_offsets=self._block_pos_offsets[:stride], - full_to_swa_mapping=pool.full_to_swa_index_mapping, - commit_lens=commit_lens, - stride=stride, - ) + if getattr(pool, "_unified_kv", False): + inject_layout = build_unified_commit_inject_layout( + req_pool_indices=batch.req_pool_indices, + prefix_lens=prefix_lens, + block_pos_offsets=self._block_pos_offsets[:stride], + commit_lens=commit_lens, + stride=stride, + ring_stride=pool.unified_swa_ring_size, + ) + else: + inject_layout = BuildCommitInjectLayout.execute( + req_pool_indices=batch.req_pool_indices, + req_to_token=self.model_runner.req_to_token_pool.req_to_token, + prefix_lens=prefix_lens, + block_pos_offsets=self._block_pos_offsets[:stride], + full_to_swa_mapping=pool.full_to_swa_index_mapping, + commit_lens=commit_lens, + stride=stride, + ) with torch.inference_mode(): self.draft_model.write_target_hidden_kv( main_hidden=hidden.reshape(-1, hidden.shape[-1]), diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index 43099c9fe5d8..222e84b5de5a 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -34,6 +34,7 @@ BuildRaggedVerifyWindow, RaggedVerifyWindow, ScatterCompactToStrided, + build_unified_commit_inject_layout, scatter_compact_to_strided_into, ) from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout @@ -304,12 +305,21 @@ def commit_hidden( if hidden is None: raise RuntimeError("DSpark verify requires target hidden states, got None.") hidden = hidden.view(bs, self.verify_num_draft_tokens, -1) + # unified_kv needs the per-token draft req slot to address the SWA ring + # (state_slot * ring + pos % ring). Verify tokens are the latest in each + # req so they always fall in the window; the commit gate (via commit_lens + # + cache_loc_2d) drops rejected tokens, so no final_pos skip is needed. + vlen = verify_window.verify_cache_loc_2d.shape[1] + state_slot = ( + batch.req_pool_indices[:bs].view(-1, 1).expand(bs, vlen).reshape(-1) + ) self.kv_injector.inject_target_hidden( target_hidden=hidden.reshape(-1, hidden.shape[-1]), cache_loc=verify_window.verify_cache_loc, cache_loc_2d=verify_window.verify_cache_loc_2d, positions=verify_window.positions_2d.reshape(-1), commit_lens=commit_lens, + state_slot=state_slot, ) def _run_ragged( @@ -632,15 +642,25 @@ def _commit_inject( torch.minimum(commit_lens, verify_lens.to(torch.int32)) * self.inject_gate_buf ) - inject_layout = BuildCommitInjectLayout.execute( - req_pool_indices=req_pool_indices, - req_to_token=ctx.resolve_req_to_token(), - prefix_lens=seq_lens[:bs], - block_pos_offsets=ctx.block_pos_offsets[: self.stride], - full_to_swa_mapping=pool.full_to_swa_index_mapping, - commit_lens=gated_commit_lens, - stride=self.stride, - ) + if pool._unified_kv: + inject_layout = build_unified_commit_inject_layout( + req_pool_indices=req_pool_indices, + prefix_lens=seq_lens[:bs], + block_pos_offsets=ctx.block_pos_offsets[: self.stride], + commit_lens=gated_commit_lens, + stride=self.stride, + ring_stride=pool.unified_swa_ring_size, + ) + else: + inject_layout = BuildCommitInjectLayout.execute( + req_pool_indices=req_pool_indices, + req_to_token=ctx.resolve_req_to_token(), + prefix_lens=seq_lens[:bs], + block_pos_offsets=ctx.block_pos_offsets[: self.stride], + full_to_swa_mapping=pool.full_to_swa_index_mapping, + commit_lens=gated_commit_lens, + stride=self.stride, + ) with torch.inference_mode(): ctx.draft_model.write_target_hidden_kv( main_hidden=self.strided_hidden[: bs * self.stride], diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index ee525f02acf0..1a617428c429 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -429,10 +429,29 @@ def _forward_prefill( ctx_lens, int(sum(batch.extend_lens)), ) + # unified_kv injects into the SWA ring keyed by (draft req slot, position); + # thread the per-token state_slot + the req's final position so the + # injector keeps only the last SWA window (older prefill tokens share a + # ring slot and would race). Cheap; only consumed under unified_kv. + state_slot = final_pos = None + from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + is_unified_kv_triton, + ) + + if is_unified_kv_triton(): + repeats = ctx_lens.to(torch.int64) + state_slot = torch.repeat_interleave( + batch.req_pool_indices.to(device=device, dtype=torch.int64), repeats + ) + final_pos = torch.repeat_interleave( + (draft_seq_lens + ctx_lens - 1).to(torch.int64), repeats + ) self._kv_injector.inject_target_hidden( target_hidden=logits_output.hidden_states, cache_loc=batch.out_cache_loc, positions=positions, + state_slot=state_slot, + final_pos=final_pos, ) logits_output.hidden_states = None diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dspark_verify_window.py b/python/sglang/srt/speculative/dspark_components/kernels/dspark_verify_window.py index 62f396d1af16..7c4eb8155be0 100644 --- a/python/sglang/srt/speculative/dspark_components/kernels/dspark_verify_window.py +++ b/python/sglang/srt/speculative/dspark_components/kernels/dspark_verify_window.py @@ -600,6 +600,37 @@ class CommitInjectLayoutResult(msgspec.Struct): positions: torch.Tensor +def build_unified_commit_inject_layout( + *, + req_pool_indices: torch.Tensor, + prefix_lens: torch.Tensor, + block_pos_offsets: torch.Tensor, + commit_lens: torch.Tensor, + stride: int, + ring_stride: int, +) -> CommitInjectLayoutResult: + """unified_kv counterpart of build_commit_inject_layout. + + Non-unified injection translates the verify tokens' full cache locs through + ``full_to_swa_mapping``; under unified_kv the SWA K lives in a ring addressed + directly by ``state_slot * ring_stride + pos % ring_stride``, so compute the + ring row here instead. Uncommitted tokens (col >= commit_len) get loc = -1 and + are skipped by the scatter. All ops are static-shape (CUDA-graph safe). + """ + bs = req_pool_indices.shape[0] + device = req_pool_indices.device + positions_2d = prefix_lens.unsqueeze(1) + block_pos_offsets[:stride] + positions = positions_2d.reshape(-1).to(torch.int64) + state_slot = ( + req_pool_indices.to(torch.int64).view(-1, 1).expand(bs, stride).reshape(-1) + ) + loc = state_slot * ring_stride + positions % ring_stride + col = torch.arange(stride, device=device).view(1, -1) + committed = (col < commit_lens.to(torch.long).view(-1, 1)).reshape(-1) + swa_loc = torch.where(committed, loc, torch.full_like(loc, -1)).to(torch.int32) + return CommitInjectLayoutResult(swa_loc=swa_loc, positions=positions) + + class BuildCommitInjectLayout: @classmethod def execute(cls, *args, **kwargs) -> CommitInjectLayoutResult: From 058e7c0002ec0b1c8ac51961e1e406ab5bf892d9 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Tue, 14 Jul 2026 10:17:38 +0000 Subject: [PATCH 03/25] hip guard --- .../dsv4/unified_kv_kernels/env_gate.py | 4 ++++ .../sglang/srt/models/deepseek_v4_dspark.py | 7 ++++-- .../dspark_components/dspark_kv_inject.py | 7 ++++-- .../dspark_components/dspark_verify.py | 24 ++++++++++++------- .../dspark_components/dspark_worker_v2.py | 9 ++++--- 5 files changed, 33 insertions(+), 18 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py index c55ba905e00e..a94808dd08f3 100644 --- a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py +++ b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py @@ -10,3 +10,7 @@ def is_unified_kv_triton() -> bool: # unified_kv_triton is only implemented on HIP (ROCm) return is_hip() and envs.SGLANG_HACK_FLASHMLA_BACKEND.get() == "unified_kv_triton" + + +def hip_unified_kv_triton_enabled() -> bool: + return is_hip() and is_unified_kv_triton() diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index 9d52b9cb774a..12880bde9849 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -11,6 +11,9 @@ from sglang.jit_kernel.dsv4 import fused_q_norm_rope, fused_rope_inplace from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config from sglang.srt.environ import envs +from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + hip_unified_kv_triton_enabled, +) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.quantization.base_config import QuantizationConfig @@ -141,7 +144,7 @@ def _store_block_kv( attn_backend, pool: DeepSeekV4TokenToKVPool, ) -> None: - if pool._unified_kv: + if hip_unified_kv_triton_enabled(): # unified_kv: SWA K lives in the shared bf16 ring (swa_kv_pool is # None). Use the unified ring write target -- get_unified_swa_loc # recomputes it from live positions for multi-step draft decode. @@ -685,7 +688,7 @@ def write_target_hidden_kv( # buffer. Same swa_loc/positions contract either way. store_kv = ( pool.set_unified_key_buffer_radix_fused_norm_rope - if pool._unified_kv + if hip_unified_kv_triton_enabled() else pool.set_swa_key_buffer_radix_fused_norm_rope ) for stage, kv in zip(self.stages, kvs): diff --git a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py index 4855eb633bbe..cb35810626bb 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py @@ -3,6 +3,9 @@ import torch from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func +from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + hip_unified_kv_triton_enabled, +) from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.speculative.dspark_components.kernels.dspark_verify_window import ( BuildCommitInjectLayout, @@ -102,7 +105,7 @@ def _inject_mla( state_slot: Optional[torch.Tensor] = None, final_pos: Optional[torch.Tensor] = None, ) -> None: - if getattr(pool, "_unified_kv", False): + if hip_unified_kv_triton_enabled(): swa_loc = self._unified_inject_loc( pool=pool, positions=positions, @@ -185,7 +188,7 @@ def inject_ragged( if hasattr(pool, "set_swa_key_buffer_radix_fused_norm_rope"): if hidden_strided.numel() == 0: return - if getattr(pool, "_unified_kv", False): + if hip_unified_kv_triton_enabled(): inject_layout = build_unified_commit_inject_layout( req_pool_indices=batch.req_pool_indices, prefix_lens=prefix_lens, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index 222e84b5de5a..879f3fd882f1 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -5,6 +5,9 @@ import msgspec import torch +from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + hip_unified_kv_triton_enabled, +) from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode @@ -305,14 +308,17 @@ def commit_hidden( if hidden is None: raise RuntimeError("DSpark verify requires target hidden states, got None.") hidden = hidden.view(bs, self.verify_num_draft_tokens, -1) - # unified_kv needs the per-token draft req slot to address the SWA ring - # (state_slot * ring + pos % ring). Verify tokens are the latest in each - # req so they always fall in the window; the commit gate (via commit_lens - # + cache_loc_2d) drops rejected tokens, so no final_pos skip is needed. - vlen = verify_window.verify_cache_loc_2d.shape[1] - state_slot = ( - batch.req_pool_indices[:bs].view(-1, 1).expand(bs, vlen).reshape(-1) - ) + state_slot = None + pool = self.kv_injector.draft_model_runner.token_to_kv_pool + if hip_unified_kv_triton_enabled(): + # unified_kv needs the per-token draft req slot to address the SWA ring + # (state_slot * ring + pos % ring). Verify tokens are the latest in each + # req so they always fall in the window; the commit gate (via commit_lens + # + cache_loc_2d) drops rejected tokens, so no final_pos skip is needed. + vlen = verify_window.verify_cache_loc_2d.shape[1] + state_slot = ( + batch.req_pool_indices[:bs].view(-1, 1).expand(bs, vlen).reshape(-1) + ) self.kv_injector.inject_target_hidden( target_hidden=hidden.reshape(-1, hidden.shape[-1]), cache_loc=verify_window.verify_cache_loc, @@ -642,7 +648,7 @@ def _commit_inject( torch.minimum(commit_lens, verify_lens.to(torch.int32)) * self.inject_gate_buf ) - if pool._unified_kv: + if hip_unified_kv_triton_enabled(): inject_layout = build_unified_commit_inject_layout( req_pool_indices=req_pool_indices, prefix_lens=seq_lens[:bs], diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 442ea5a5bcbf..10519c64edeb 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -5,6 +5,9 @@ import torch from sglang.srt.environ import envs +from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + hip_unified_kv_triton_enabled, +) from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker @@ -426,11 +429,7 @@ def _forward_prefill( # injector keeps only the last SWA window (older prefill tokens share a # ring slot and would race). Cheap; only consumed under unified_kv. state_slot = final_pos = None - from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( - is_unified_kv_triton, - ) - - if is_unified_kv_triton(): + if hip_unified_kv_triton_enabled(): repeats = ctx_lens.to(torch.int64) state_slot = torch.repeat_interleave( batch.req_pool_indices.to(device=device, dtype=torch.int64), repeats From 909db6a417183fa837d5634c95c4c1abf7219ed7 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Tue, 14 Jul 2026 10:25:05 +0000 Subject: [PATCH 04/25] x --- .../layers/attention/dsv4/unified_kv_kernels/env_gate.py | 4 ---- python/sglang/srt/models/deepseek_v4_dspark.py | 6 +++--- .../srt/speculative/dspark_components/dspark_kv_inject.py | 6 +++--- .../srt/speculative/dspark_components/dspark_verify.py | 6 +++--- .../srt/speculative/dspark_components/dspark_worker_v2.py | 4 ++-- 5 files changed, 11 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py index a94808dd08f3..c55ba905e00e 100644 --- a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py +++ b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py @@ -10,7 +10,3 @@ def is_unified_kv_triton() -> bool: # unified_kv_triton is only implemented on HIP (ROCm) return is_hip() and envs.SGLANG_HACK_FLASHMLA_BACKEND.get() == "unified_kv_triton" - - -def hip_unified_kv_triton_enabled() -> bool: - return is_hip() and is_unified_kv_triton() diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index 12880bde9849..9b122a4f2a24 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -12,7 +12,7 @@ from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config from sglang.srt.environ import envs from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( - hip_unified_kv_triton_enabled, + is_unified_kv_triton, ) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.logits_processor import LogitsProcessorOutput @@ -144,7 +144,7 @@ def _store_block_kv( attn_backend, pool: DeepSeekV4TokenToKVPool, ) -> None: - if hip_unified_kv_triton_enabled(): + if is_unified_kv_triton(): # unified_kv: SWA K lives in the shared bf16 ring (swa_kv_pool is # None). Use the unified ring write target -- get_unified_swa_loc # recomputes it from live positions for multi-step draft decode. @@ -688,7 +688,7 @@ def write_target_hidden_kv( # buffer. Same swa_loc/positions contract either way. store_kv = ( pool.set_unified_key_buffer_radix_fused_norm_rope - if hip_unified_kv_triton_enabled() + if is_unified_kv_triton() else pool.set_swa_key_buffer_radix_fused_norm_rope ) for stage, kv in zip(self.stages, kvs): diff --git a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py index cb35810626bb..27b23e182bdf 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py @@ -4,7 +4,7 @@ from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( - hip_unified_kv_triton_enabled, + is_unified_kv_triton, ) from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.speculative.dspark_components.kernels.dspark_verify_window import ( @@ -105,7 +105,7 @@ def _inject_mla( state_slot: Optional[torch.Tensor] = None, final_pos: Optional[torch.Tensor] = None, ) -> None: - if hip_unified_kv_triton_enabled(): + if is_unified_kv_triton(): swa_loc = self._unified_inject_loc( pool=pool, positions=positions, @@ -188,7 +188,7 @@ def inject_ragged( if hasattr(pool, "set_swa_key_buffer_radix_fused_norm_rope"): if hidden_strided.numel() == 0: return - if hip_unified_kv_triton_enabled(): + if is_unified_kv_triton(): inject_layout = build_unified_commit_inject_layout( req_pool_indices=batch.req_pool_indices, prefix_lens=prefix_lens, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index 879f3fd882f1..9173143e7406 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -6,7 +6,7 @@ import torch from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( - hip_unified_kv_triton_enabled, + is_unified_kv_triton, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.schedule_batch import ScheduleBatch @@ -310,7 +310,7 @@ def commit_hidden( hidden = hidden.view(bs, self.verify_num_draft_tokens, -1) state_slot = None pool = self.kv_injector.draft_model_runner.token_to_kv_pool - if hip_unified_kv_triton_enabled(): + if is_unified_kv_triton(): # unified_kv needs the per-token draft req slot to address the SWA ring # (state_slot * ring + pos % ring). Verify tokens are the latest in each # req so they always fall in the window; the commit gate (via commit_lens @@ -648,7 +648,7 @@ def _commit_inject( torch.minimum(commit_lens, verify_lens.to(torch.int32)) * self.inject_gate_buf ) - if hip_unified_kv_triton_enabled(): + if is_unified_kv_triton(): inject_layout = build_unified_commit_inject_layout( req_pool_indices=req_pool_indices, prefix_lens=seq_lens[:bs], diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 10519c64edeb..abaa0f70c557 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -6,7 +6,7 @@ from sglang.srt.environ import envs from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( - hip_unified_kv_triton_enabled, + is_unified_kv_triton, ) from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult @@ -429,7 +429,7 @@ def _forward_prefill( # injector keeps only the last SWA window (older prefill tokens share a # ring slot and would race). Cheap; only consumed under unified_kv. state_slot = final_pos = None - if hip_unified_kv_triton_enabled(): + if is_unified_kv_triton(): repeats = ctx_lens.to(torch.int64) state_slot = torch.repeat_interleave( batch.req_pool_indices.to(device=device, dtype=torch.int64), repeats From 888db411a551d1aee387a1a45e289ca4430bde35 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Tue, 14 Jul 2026 12:37:38 +0000 Subject: [PATCH 05/25] ci --- .../spec/dspark/test_dspark_kernel_parity.py | 3 +- .../test_ragged_verify_backend_capability.py | 41 ++++++++++++------- 2 files changed, 29 insertions(+), 15 deletions(-) diff --git a/test/registered/spec/dspark/test_dspark_kernel_parity.py b/test/registered/spec/dspark/test_dspark_kernel_parity.py index 9faa2f7a612b..f606400fe67c 100644 --- a/test/registered/spec/dspark/test_dspark_kernel_parity.py +++ b/test/registered/spec/dspark/test_dspark_kernel_parity.py @@ -25,10 +25,11 @@ dspark_verify_window, ) from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small") +register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd") DEVICE = torch.device("cuda") VOCAB = 129280 diff --git a/test/registered/spec/dspark/test_ragged_verify_backend_capability.py b/test/registered/spec/dspark/test_ragged_verify_backend_capability.py index b0257ae28bc8..fa377bd15d59 100644 --- a/test/registered/spec/dspark/test_ragged_verify_backend_capability.py +++ b/test/registered/spec/dspark/test_ragged_verify_backend_capability.py @@ -6,10 +6,12 @@ import unittest -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.srt.utils import is_hip +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") +register_amd_ci(est_time=15, suite="stage-b-test-1-gpu-small-amd") class TestRaggedVerifyGraphCapability(CustomTestCase): @@ -22,19 +24,30 @@ def test_ragged_implementing_backends_declare_the_flag(self): """Every backend with a ragged-verify metadata path must opt in; a dropped flag silently disables ragged graphs for that backend (the runner falls back to eager with no other test going red).""" - from sglang.srt.layers.attention.deepseek_v4_backend import ( - DeepseekV4AttnBackend, - ) - from sglang.srt.layers.attention.flashattention_backend import ( - FlashAttentionBackend, - ) - from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend - - for backend in ( - TRTLLMHAAttnBackend, - DeepseekV4AttnBackend, - FlashAttentionBackend, - ): + if is_hip(): + from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import ( + DeepseekV4HipRadixBackend, + ) + + backends = (DeepseekV4HipRadixBackend,) + else: + from sglang.srt.layers.attention.deepseek_v4_backend import ( + DeepseekV4AttnBackend, + ) + from sglang.srt.layers.attention.flashattention_backend import ( + FlashAttentionBackend, + ) + from sglang.srt.layers.attention.trtllm_mha_backend import ( + TRTLLMHAAttnBackend, + ) + + backends = ( + TRTLLMHAAttnBackend, + DeepseekV4AttnBackend, + FlashAttentionBackend, + ) + + for backend in backends: with self.subTest(backend=backend.__name__): self.assertTrue(backend.supports_ragged_verify_graph) From c792ee9d11f0abd19c5b73d2610f4aa64aab1512 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Tue, 14 Jul 2026 12:39:08 +0000 Subject: [PATCH 06/25] Revert "ci" This reverts commit 888db411a551d1aee387a1a45e289ca4430bde35. --- .../spec/dspark/test_dspark_kernel_parity.py | 3 +- .../test_ragged_verify_backend_capability.py | 41 +++++++------------ 2 files changed, 15 insertions(+), 29 deletions(-) diff --git a/test/registered/spec/dspark/test_dspark_kernel_parity.py b/test/registered/spec/dspark/test_dspark_kernel_parity.py index f606400fe67c..9faa2f7a612b 100644 --- a/test/registered/spec/dspark/test_dspark_kernel_parity.py +++ b/test/registered/spec/dspark/test_dspark_kernel_parity.py @@ -25,11 +25,10 @@ dspark_verify_window, ) from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small") -register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd") DEVICE = torch.device("cuda") VOCAB = 129280 diff --git a/test/registered/spec/dspark/test_ragged_verify_backend_capability.py b/test/registered/spec/dspark/test_ragged_verify_backend_capability.py index fa377bd15d59..b0257ae28bc8 100644 --- a/test/registered/spec/dspark/test_ragged_verify_backend_capability.py +++ b/test/registered/spec/dspark/test_ragged_verify_backend_capability.py @@ -6,12 +6,10 @@ import unittest -from sglang.srt.utils import is_hip -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") -register_amd_ci(est_time=15, suite="stage-b-test-1-gpu-small-amd") class TestRaggedVerifyGraphCapability(CustomTestCase): @@ -24,30 +22,19 @@ def test_ragged_implementing_backends_declare_the_flag(self): """Every backend with a ragged-verify metadata path must opt in; a dropped flag silently disables ragged graphs for that backend (the runner falls back to eager with no other test going red).""" - if is_hip(): - from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import ( - DeepseekV4HipRadixBackend, - ) - - backends = (DeepseekV4HipRadixBackend,) - else: - from sglang.srt.layers.attention.deepseek_v4_backend import ( - DeepseekV4AttnBackend, - ) - from sglang.srt.layers.attention.flashattention_backend import ( - FlashAttentionBackend, - ) - from sglang.srt.layers.attention.trtllm_mha_backend import ( - TRTLLMHAAttnBackend, - ) - - backends = ( - TRTLLMHAAttnBackend, - DeepseekV4AttnBackend, - FlashAttentionBackend, - ) - - for backend in backends: + from sglang.srt.layers.attention.deepseek_v4_backend import ( + DeepseekV4AttnBackend, + ) + from sglang.srt.layers.attention.flashattention_backend import ( + FlashAttentionBackend, + ) + from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend + + for backend in ( + TRTLLMHAAttnBackend, + DeepseekV4AttnBackend, + FlashAttentionBackend, + ): with self.subTest(backend=backend.__name__): self.assertTrue(backend.supports_ragged_verify_graph) From bbd7f3151521f0376cb091aa6b83b95429170b62 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Tue, 14 Jul 2026 13:30:17 +0000 Subject: [PATCH 07/25] ci test --- .../amd/test_deepseek_v4_pro_fp4_dspark.py | 157 ++++++++++++++++++ 1 file changed, 157 insertions(+) create mode 100644 test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py diff --git a/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py b/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py new file mode 100644 index 000000000000..5f46b4cfee12 --- /dev/null +++ b/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py @@ -0,0 +1,157 @@ +"""MI35x DeepSeek-V4-Pro-DSpark unified_kv GSM8K accuracy test (8-GPU). + +Runs the production AMD DSpark static configuration with the HIP dsv4 backend and +SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton. The test uses the full GSM8K set +to catch regressions in unified-KV target-hidden injection, verify metadata, and +DSpark acceptance. + +Registry: nightly-amd-8-gpu-mi35x-deepseek-v4-pro-dspark suite +""" + +import os +import unittest +from types import SimpleNamespace + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + is_in_ci, + popen_launch_server, + write_github_step_summary, +) + +register_amd_ci( + est_time=7200, suite="nightly-amd-8-gpu-mi35x-deepseek-v4-pro-dspark", nightly=True +) + +DEEPSEEK_V4_DSPARK_MODEL_PATH = os.environ.get( + "DEEPSEEK_V4_DSPARK_MODEL_PATH", "deepseek-ai/DeepSeek-V4-Pro-DSpark" +) +SERVER_LAUNCH_TIMEOUT = 5400 +GSM8K_ACCURACY_THRESHOLD = 0.92 +AVG_SPEC_ACCEPT_LENGTH_THRESHOLD = 4.0 + + +class TestDeepseekV4DSparkUnifiedKVGSM8K(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEEPSEEK_V4_DSPARK_MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + env = os.environ.copy() + env.update( + { + "SGLANG_DEFAULT_THINKING": "1", + "SGLANG_DSV4_REASONING_EFFORT": "max", + "SGLANG_OPT_DEEPGEMM_HC_PRENORM": "false", + "SGLANG_USE_AITER": "1", + "SGLANG_USE_ROCM700A": "0", + "SGLANG_OPT_USE_FUSED_COMPRESS": "true", + "SGLANG_HACK_FLASHMLA_BACKEND": "unified_kv_triton", + "SGLANG_RAGGED_VERIFY_MODE": "static", + "SGLANG_DSPARK_ENABLE_SPS_ONLINE_PROFILE": "0", + "SGLANG_OPT_FP8_WO_A_GEMM": "false", + "SGLANG_OPT_USE_JIT_INDEXER_METADATA": "false", + "SGLANG_OPT_USE_TOPK_V2": "false", + "SGLANG_OPT_USE_AITER_INDEXER": "true", + "SGLANG_OPT_USE_TILELANG_INDEXER": "false", + "SGLANG_OPT_USE_TILELANG_MHC_PRE": "false", + "SGLANG_OPT_USE_TILELANG_MHC_POST": "false", + "SGLANG_FP8_PAGED_MQA_LOGITS_TORCH": "1", + "SGLANG_OPT_USE_FUSED_COMPRESS_TRITON": "true", + "SGLANG_OPT_USE_MULTI_STREAM_OVERLAP": "false", + "SGLANG_ROCM_USE_MULTI_STREAM": "false", + "AITER_BF16_FP8_MOE_BOUND": "0", + "SGLANG_EAGER_INPUT_NO_COPY": "true", + "SGLANG_SHARED_EXPERT_TP1": "1", + "SGLANG_DP_SHARED_EXPERT_LOCAL": "1", + "SGLANG_DP_USE_GATHERV": "1", + "SGLANG_DP_USE_REDUCE_SCATTER": "1", + "GPU_MAX_HW_QUEUES": "5", + } + ) + other_args = [ + "--trust-remote-code", + "--tp", + "8", + "--dp", + "8", + "--enable-dp-attention", + "--enable-dp-lm-head", + "--enable-prefill-delayer", + "--disable-radix-cache", + "--attention-backend", + "dsv4", + "--page-size", + "256", + "--mem-fraction-static", + "0.9", + "--swa-full-tokens-ratio", + "0.15", + "--disable-shared-experts-fusion", + "--tool-call-parser", + "deepseekv4", + "--reasoning-parser", + "deepseek-v4", + "--kv-cache-dtype", + "fp8_e4m3", + "--chunked-prefill-size", + "65536", + "--cuda-graph-max-bs", + "512", + "--max-running-requests", + "512", + "--speculative-algorithm", + "DSPARK", + "--speculative-dspark-block-size", + "5", + ] + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=other_args, + env=env, + ) + + @classmethod + def tearDownClass(cls): + if getattr(cls, "process", None) is not None: + kill_process_tree(cls.process.pid) + + def test_full_gsm8k_unified_kv_dspark_static(self): + requests.get(self.base_url + "/flush_cache") + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=1319, + parallel=512, + max_new_tokens=512, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval_few_shot_gsm8k(args) + print(f"{metrics=}") + + server_info = requests.get(self.base_url + "/server_info") + avg_spec_accept_length = server_info.json()["internal_states"][0][ + "avg_spec_accept_length" + ] + print(f"{avg_spec_accept_length=}") + + if is_in_ci(): + write_github_step_summary( + "### test_gsm8k (deepseek-v4-pro-dspark unified_kv static MI35x)\n" + f"accuracy={metrics['accuracy']:.3f}\n" + f"avg_spec_accept_length={avg_spec_accept_length:.2f}\n" + ) + self.assertGreater(metrics["accuracy"], GSM8K_ACCURACY_THRESHOLD) + self.assertGreater(avg_spec_accept_length, AVG_SPEC_ACCEPT_LENGTH_THRESHOLD) + + +if __name__ == "__main__": + unittest.main() From ac2470d92a145deb4ac3e4b901d8d83378a74404 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Tue, 14 Jul 2026 13:40:32 +0000 Subject: [PATCH 08/25] Revert "Delete CUTLASS FP8 blockwise for SM90 and SM100, move SM120 to JIT and add SwapAB (#30438)" This reverts commit 7431f35fd8a93c4a8a193d372f7697b55cf9b6d0. --- .../fp8_blockwise_scaled_mm_entry.cuh | 25 - .../fp8_blockwise_scaled_mm_sm120.cuh | 502 ----------------- .../sglang/jit_kernel/fp8_blockwise_gemm.py | 93 ---- .../jit_kernel/include/sgl_kernel/utils.cuh | 16 - .../srt/layers/quantization/fp8_utils.py | 29 +- scripts/ci/cuda/warmup_deep_gemm.py | 2 +- sgl-kernel/CMakeLists.txt | 1 + .../benchmark/bench_fp8_blockwise_gemm.py | 237 ++++++++ sgl-kernel/csrc/common_extension.cc | 5 + .../gemm/fp8_blockwise_gemm_sm90_dispatch.cuh | 197 +++++++ .../csrc/gemm/fp8_blockwise_gemm_kernel.cu | 522 ++++++++++++++++++ sgl-kernel/include/sgl_kernel_ops.h | 6 + sgl-kernel/python/sgl_kernel/__init__.py | 2 + sgl-kernel/python/sgl_kernel/gemm.py | 10 + .../tests}/test_fp8_blockwise_gemm.py | 38 +- .../jit/benchmark/bench_fp8_blockwise_gemm.py | 103 ---- .../quant/test_fp8_blockwise_row_padding.py | 136 +++++ 17 files changed, 1159 insertions(+), 765 deletions(-) delete mode 100644 python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh delete mode 100644 python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh delete mode 100644 python/sglang/jit_kernel/fp8_blockwise_gemm.py create mode 100644 sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py create mode 100644 sgl-kernel/csrc/cutlass_extensions/gemm/fp8_blockwise_gemm_sm90_dispatch.cuh create mode 100644 sgl-kernel/csrc/gemm/fp8_blockwise_gemm_kernel.cu rename {test/registered/jit => sgl-kernel/tests}/test_fp8_blockwise_gemm.py (69%) delete mode 100644 test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py create mode 100644 test/registered/quant/test_fp8_blockwise_row_padding.py diff --git a/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh b/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh deleted file mode 100644 index cdbecf8a2895..000000000000 --- a/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh +++ /dev/null @@ -1,25 +0,0 @@ -/* Copyright 2026 SGLang Team. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "fp8_blockwise_scaled_mm_sm120.cuh" - -void fp8_blockwise_scaled_mm( - tvm::ffi::TensorView out, - tvm::ffi::TensorView mat_a, - tvm::ffi::TensorView mat_b, - tvm::ffi::TensorView scales_a, - tvm::ffi::TensorView scales_b) { - fp8_blockwise_scaled_mm_sm120(out, mat_a, mat_b, scales_a, scales_b); -} diff --git a/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh b/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh deleted file mode 100644 index 8f802d92efbe..000000000000 --- a/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh +++ /dev/null @@ -1,502 +0,0 @@ -/* Copyright 2026 SGLang Team. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#pragma once - -#include -#include - -#include -#include - -#include -#include -#include - -using namespace host; - -// clang-format off -#include "cutlass/cutlass.h" -#include "cutlass/detail/blockwise_scale_layout.hpp" -#include "cutlass/gemm/collective/collective_builder.hpp" -#include "cutlass/epilogue/collective/collective_builder.hpp" -#include "cutlass/gemm/device/gemm_universal_adapter.h" -#include "cutlass/gemm/kernel/gemm_universal.hpp" -#include "cutlass/gemm/dispatch_policy.hpp" -#include "cutlass/util/packed_stride.hpp" -// clang-format on - -#define CUTLASS_CHECK(status) \ - { \ - cutlass::Status error = status; \ - RuntimeCheck(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \ - } - -using namespace cute; - -#if defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED) - -template < - typename OutType, - typename MmaTileShape, - typename PerSmTileShape, - typename EpilogueTileShape, - typename ScalesPerTile, - int TileSizeM_ = 128, - class ClusterShape = Shape<_1, _1, _1>> -void launch_sm120_fp8_blockwise_scaled_mm( - tvm::ffi::TensorView out, - tvm::ffi::TensorView a, - tvm::ffi::TensorView b, - tvm::ffi::TensorView scales_a, - tvm::ffi::TensorView scales_b, - cudaStream_t stream) { - using ElementBlockScale = float; - - // A matrix configuration - using ElementA = cutlass::float_e4m3_t; // Element type for A matrix operand - using LayoutATag = cutlass::layout::RowMajor; // Layout type for A matrix operand - constexpr int AlignmentA = - 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of A matrix in units of - // elements (up to 16 bytes) - - // B matrix configuration - using ElementB = cutlass::float_e4m3_t; // Element type for B matrix operand - using LayoutBTag = cutlass::layout::ColumnMajor; // Layout type for B matrix operand - constexpr int AlignmentB = - 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of B matrix in units of - // elements (up to 16 bytes) - - // C/D matrix configuration - using ElementD = OutType; // Element type for D matrix operand - using ElementC = void; // Element type for C matrix operand - using LayoutCTag = cutlass::layout::RowMajor; // Layout type for C matrix operand - using LayoutDTag = cutlass::layout::RowMajor; // Layout type for D matrix operand - constexpr int AlignmentD = - 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of C matrix in units of - // elements (up to 16 bytes) - constexpr int AlignmentC = - AlignmentD; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes) - - // Kernel functional config - using ElementAccumulator = float; // Element type for internal accumulation - using ArchTag = cutlass::arch::Sm120; // Tag indicating the minimum SM that supports the intended feature - using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag - changed from OpClassBlockScaledTensorOp - - static constexpr int ScaleMsPerTile = size<0>(ScalesPerTile{}); - static constexpr int ScaleGranularityM = size<0>(MmaTileShape{}) / ScaleMsPerTile; - static constexpr int ScaleGranularityN = size<1>(MmaTileShape{}) / size<1>(ScalesPerTile{}); - static constexpr int ScaleGranularityK = size<2>(MmaTileShape{}) / size<2>(ScalesPerTile{}); - - using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig< - ScaleGranularityM, - ScaleGranularityN, - ScaleGranularityK, - cute::UMMA::Major::MN, - cute::UMMA::Major::K>; - // FP8 Block-wise scaling configuration - using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); // Layout type for SFA matrix operand - using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); // Layout type for SFB matrix operand - - constexpr bool kCanUsePingpong = (64 % ScaleGranularityM == 0); - - int m = a.size(0); - int k = a.size(1); - int n = b.size(1); - - auto a_ptr = static_cast(a.data_ptr()); - auto b_ptr = static_cast(b.data_ptr()); - auto c_ptr = static_cast(out.data_ptr()); - - auto scales_a_ptr = static_cast(scales_a.data_ptr()); - auto scales_b_ptr = static_cast(scales_b.data_ptr()); - - LayoutSFA layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1)); - LayoutSFB layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1)); - - auto run_gemm = [&](auto tag) -> cutlass::Status { - using GemmKernel = decltype(tag); - using Gemm = cutlass::gemm::device::GemmUniversalAdapter; - Gemm gemm_op; - - using StrideA = typename GemmKernel::StrideA; - using StrideB = typename GemmKernel::StrideB; - using StrideC = typename GemmKernel::StrideD; - - StrideA stride_a = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1)); - StrideB stride_b = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1)); - StrideC stride_c = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1)); - - typename GemmKernel::MainloopArguments mainloop_args{ - a_ptr, stride_a, b_ptr, stride_b, scales_a_ptr, layout_SFA, scales_b_ptr, layout_SFB}; - - typename GemmKernel::EpilogueArguments epilogue_args{{}, c_ptr, stride_c, c_ptr, stride_c}; - epilogue_args.thread.alpha = 1.0f; - - typename Gemm::Arguments args = { - cutlass::gemm::GemmUniversalMode::kGemm, - {m, n, k, 1}, - mainloop_args, - epilogue_args, - }; - - auto can_implement = gemm_op.can_implement(args); - if (can_implement != cutlass::Status::kSuccess) { - return can_implement; - } - - size_t workspace_size = gemm_op.get_workspace_size(args); - auto workspace_tensor = alloc_workspace_tensor(workspace_size, a.device()); - void* workspace = (workspace_size == 0) ? nullptr : workspace_tensor.data_ptr(); - - auto init_status = gemm_op.initialize(args, workspace, stream); - if (init_status != cutlass::Status::kSuccess) { - return init_status; - } - - return gemm_op.run(stream); - }; - - using CooperativeCollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - PerSmTileShape, - ClusterShape, - cutlass::epilogue::collective::EpilogueTileAuto, - ElementAccumulator, - ElementAccumulator, - ElementC, - LayoutCTag, - AlignmentC, - ElementD, - LayoutDTag, - AlignmentD, - cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; - - using CooperativeStageCount = cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename CooperativeCollectiveEpilogue::SharedStorage))>; - - using CooperativeCollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - ElementA, - cute::tuple, - AlignmentA, - ElementB, - cute::tuple, - AlignmentB, - ElementAccumulator, - MmaTileShape, - ClusterShape, - CooperativeStageCount, - cutlass::gemm::KernelScheduleSm120Blockwise>::CollectiveOp; - - using CooperativeGemmKernelStreamK = cutlass::gemm::kernel::GemmUniversal< - Shape, - CooperativeCollectiveMainloop, - CooperativeCollectiveEpilogue, - cutlass::gemm::StreamKScheduler>; - using CooperativeGemmKernelVoid = cutlass::gemm::kernel:: - GemmUniversal, CooperativeCollectiveMainloop, CooperativeCollectiveEpilogue, void>; - - auto run_cooperative = [&]() -> cutlass::Status { - static const uint32_t kNumSM = host::runtime::get_sm_count(a.device().device_id); - constexpr int kTileM = size<0>(MmaTileShape{}); - constexpr int kTileN = size<1>(MmaTileShape{}); - uint64_t tiles = static_cast((m + kTileM - 1) / kTileM) * ((n + kTileN - 1) / kTileN); - uint32_t last_wave = static_cast(tiles % kNumSM); - if (last_wave == 0) last_wave = kNumSM; - float waste = 1.0f - static_cast(last_wave) / static_cast(kNumSM); - return (waste > 0.5f) ? run_gemm(CooperativeGemmKernelStreamK{}) : run_gemm(CooperativeGemmKernelVoid{}); - }; - - cutlass::Status status = cutlass::Status::kSuccess; - if constexpr (kCanUsePingpong) { - using PingpongMmaTileShape_MNK = Shape<_64, _128, _128>; - using PingpongCollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - PerSmTileShape, - ClusterShape, - cutlass::epilogue::collective::EpilogueTileAuto, - ElementAccumulator, - ElementAccumulator, - ElementC, - LayoutCTag, - AlignmentC, - ElementD, - LayoutDTag, - AlignmentD, - cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; - - using PingpongStageCount = cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename PingpongCollectiveEpilogue::SharedStorage))>; - - using PingpongCollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - ElementA, - cute::tuple, - AlignmentA, - ElementB, - cute::tuple, - AlignmentB, - ElementAccumulator, - PingpongMmaTileShape_MNK, - ClusterShape, - PingpongStageCount, - cutlass::gemm::KernelTmaWarpSpecializedBlockwisePingpongSm120>::CollectiveOp; - - using PingpongGemmKernel = cutlass::gemm::kernel:: - GemmUniversal, PingpongCollectiveMainloop, PingpongCollectiveEpilogue, void>; - - if (m <= 64) { - status = run_gemm(PingpongGemmKernel{}); - if (status != cutlass::Status::kSuccess) { - status = run_cooperative(); - } - } else { - status = run_cooperative(); - } - } else { - status = run_cooperative(); - } - - CUTLASS_CHECK(status); -} - -// Transposed GEMM D^T = Wgemm(weight, activation): puts tokens on the N axis. -template < - typename OutType, - typename MmaTileShape, - typename PerSmTileShape, - typename EpilogueTileShape, - typename ScalesPerTile, - class ClusterShape = Shape<_1, _1, _1>> -void launch_sm120_fp8_blockwise_scaled_mm_swapab( - tvm::ffi::TensorView out, - tvm::ffi::TensorView a, - tvm::ffi::TensorView b, - tvm::ffi::TensorView scales_a, - tvm::ffi::TensorView scales_b, - cudaStream_t stream) { - using ElementBlockScale = float; - - using ElementA = cutlass::float_e4m3_t; // A' = weight - using LayoutATag = cutlass::layout::RowMajor; // weight [N, K] is row-major - constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; - - using ElementB = cutlass::float_e4m3_t; // B' = activation - using LayoutBTag = cutlass::layout::ColumnMajor; // activation as [K, M] column-major - constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; - - using ElementD = OutType; - using ElementC = void; - using LayoutCTag = cutlass::layout::ColumnMajor; // D' = out^T is column-major - using LayoutDTag = cutlass::layout::ColumnMajor; - constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; - constexpr int AlignmentC = AlignmentD; - - using ElementAccumulator = float; - using ArchTag = cutlass::arch::Sm120; - using OperatorClass = cutlass::arch::OpClassTensorOp; - - static constexpr int ScaleMsPerTile = size<0>(ScalesPerTile{}); - static constexpr int ScaleGranularityM = size<0>(MmaTileShape{}) / ScaleMsPerTile; - static constexpr int ScaleGranularityN = size<1>(MmaTileShape{}) / size<1>(ScalesPerTile{}); - static constexpr int ScaleGranularityK = size<2>(MmaTileShape{}) / size<2>(ScalesPerTile{}); - - // Operands are swapped, so the scale majors swap relative to the non-swap path: - // SFA (weight) is K-major; SFB (per-token activation) is MN-major. - using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig< - ScaleGranularityM, - ScaleGranularityN, - ScaleGranularityK, - cute::UMMA::Major::K, - cute::UMMA::Major::MN>; - using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); - using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); - - int m = a.size(0); // original tokens -> swapped N' - int k = a.size(1); - int n = b.size(1); // original weight cols -> swapped M' - - auto weight_ptr = static_cast(b.data_ptr()); - auto act_ptr = static_cast(a.data_ptr()); - auto c_ptr = static_cast(out.data_ptr()); - auto weight_scale_ptr = static_cast(scales_b.data_ptr()); - auto act_scale_ptr = static_cast(scales_a.data_ptr()); - - // Swapped problem shape (M', N', K) = (n, m, k). - LayoutSFA layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(make_shape(n, m, k, 1)); - LayoutSFB layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1)); - - auto run_gemm = [&](auto tag) -> cutlass::Status { - using GemmKernel = decltype(tag); - using Gemm = cutlass::gemm::device::GemmUniversalAdapter; - Gemm gemm_op; - - using StrideA = typename GemmKernel::StrideA; - using StrideB = typename GemmKernel::StrideB; - using StrideC = typename GemmKernel::StrideD; - - StrideA stride_a = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(n, k, 1)); - StrideB stride_b = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(m, k, 1)); - StrideC stride_c = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(n, m, 1)); - - typename GemmKernel::MainloopArguments mainloop_args{ - weight_ptr, stride_a, act_ptr, stride_b, weight_scale_ptr, layout_SFA, act_scale_ptr, layout_SFB}; - - typename GemmKernel::EpilogueArguments epilogue_args{{}, c_ptr, stride_c, c_ptr, stride_c}; - epilogue_args.thread.alpha = 1.0f; - - typename Gemm::Arguments args = { - cutlass::gemm::GemmUniversalMode::kGemm, - {n, m, k, 1}, - mainloop_args, - epilogue_args, - }; - - auto can_implement = gemm_op.can_implement(args); - if (can_implement != cutlass::Status::kSuccess) { - return can_implement; - } - - size_t workspace_size = gemm_op.get_workspace_size(args); - auto workspace_tensor = alloc_workspace_tensor(workspace_size, a.device()); - void* workspace = (workspace_size == 0) ? nullptr : workspace_tensor.data_ptr(); - - auto init_status = gemm_op.initialize(args, workspace, stream); - if (init_status != cutlass::Status::kSuccess) { - return init_status; - } - - return gemm_op.run(stream); - }; - - using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - PerSmTileShape, - ClusterShape, - cutlass::epilogue::collective::EpilogueTileAuto, - ElementAccumulator, - ElementAccumulator, - ElementC, - LayoutCTag, - AlignmentC, - ElementD, - LayoutDTag, - AlignmentD, - cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; - - using StageCount = cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename CollectiveEpilogue::SharedStorage))>; - - using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - ElementA, - cute::tuple, - AlignmentA, - ElementB, - cute::tuple, - AlignmentB, - ElementAccumulator, - MmaTileShape, - ClusterShape, - StageCount, - cutlass::gemm::KernelScheduleSm120Blockwise>::CollectiveOp; - - using GemmKernel = - cutlass::gemm::kernel::GemmUniversal, CollectiveMainloop, CollectiveEpilogue, void>; - - CUTLASS_CHECK(run_gemm(GemmKernel{})); -} - -// swapAB (tile N=32) beats the non-swap 128x128 path for M<=64 or M%4!=0 -// (cold-L2 CUPTI benchmarks, up to ~1.2x); tile N=16 is unsupported by the -// SM120 blockwise collective (needs EPI_TILE_N=32 | CTA_N and B LDSM N>=32). -template -void sm120_fp8_blockwise_dispatch_shape( - tvm::ffi::TensorView out, - tvm::ffi::TensorView a, - tvm::ffi::TensorView b, - tvm::ffi::TensorView scales_a, - tvm::ffi::TensorView scales_b, - cudaStream_t stream) { - const int m = a.size(0); - using EpilogueTileShape = Shape<_128, _64>; - if (m <= 64 || (m % 4 != 0)) { - launch_sm120_fp8_blockwise_scaled_mm_swapab< - OutType, - Shape<_128, _32, _128>, - Shape<_128, _32, _128>, - EpilogueTileShape, - Shape<_1, _32, _1>>(out, a, b, scales_a, scales_b, stream); - return; - } - - using MmaTileShape = Shape<_128, _128, _128>; - using PerSmTileShape = Shape<_128, _128, _128>; - using ScalesPerTile = Shape<_128, _1, _1>; - launch_sm120_fp8_blockwise_scaled_mm( - out, a, b, scales_a, scales_b, stream); -} - -inline void fp8_blockwise_scaled_mm_sm120( - tvm::ffi::TensorView out, - tvm::ffi::TensorView mat_a, - tvm::ffi::TensorView mat_b, - tvm::ffi::TensorView scales_a, - tvm::ffi::TensorView scales_b) { - RuntimeCheck(mat_a.device().device_type == kDLCUDA, "mat_a must be a CUDA tensor"); - RuntimeCheck(mat_b.device().device_type == kDLCUDA, "mat_b must be a CUDA tensor"); - - RuntimeCheck(mat_a.dim() == 2, "mat_a must be a 2D tensor"); - RuntimeCheck(mat_b.dim() == 2, "mat_b must be a 2D tensor"); - RuntimeCheck(mat_a.stride(1) == 1, "mat_a must be a row major tensor"); - RuntimeCheck(mat_b.stride(0) == 1, "mat_b must be a column major tensor"); - RuntimeCheck(mat_a.size(1) == mat_b.size(0), "mat_a and mat_b shapes cannot be multiplied"); - - RuntimeCheck( - (mat_a.size(1) * (mat_a.dtype().bits / 8)) % 16 == 0, "mat_a must be multiple of 16 bytes for memory alignment"); - RuntimeCheck( - (mat_b.size(0) * (mat_b.dtype().bits / 8)) % 16 == 0, "mat_b must be multiple of 16 bytes for memory alignment"); - RuntimeCheck(host::is_type(mat_a.dtype()), "mat_a must be Float8_e4m3fn"); - RuntimeCheck(host::is_type(mat_b.dtype()), "mat_b must be Float8_e4m3fn"); - - RuntimeCheck(mat_a.size(0) == scales_a.size(0), "size of scales_a is not matched"); - RuntimeCheck(mat_a.size(1) / 128 == scales_a.size(1), "size of scales_a is not matched"); - RuntimeCheck(mat_b.size(0) / 128 == scales_b.size(0), "size of scales_b is not matched"); - RuntimeCheck(mat_b.size(1) / 128 == scales_b.size(1), "size of scales_b is not matched"); - RuntimeCheck(host::is_type(scales_a.dtype()), "scales_a must be Float32"); - RuntimeCheck(host::is_type(scales_b.dtype()), "scales_b must be Float32"); - - RuntimeCheck( - (out.size(1) * (out.dtype().bits / 8)) % 16 == 0, "out must be multiple of 16 bytes for memory alignment"); - - const cudaStream_t stream = LaunchKernel::resolve_device(mat_a.device()); - - if (host::is_type(out.dtype())) { - sm120_fp8_blockwise_dispatch_shape(out, mat_a, mat_b, scales_a, scales_b, stream); - } else if (host::is_type(out.dtype())) { - sm120_fp8_blockwise_dispatch_shape(out, mat_a, mat_b, scales_a, scales_b, stream); - } else { - Panic("out_dtype must be Half or BFloat16"); - } -} - -#endif // defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED) diff --git a/python/sglang/jit_kernel/fp8_blockwise_gemm.py b/python/sglang/jit_kernel/fp8_blockwise_gemm.py deleted file mode 100644 index 49b4c9606e25..000000000000 --- a/python/sglang/jit_kernel/fp8_blockwise_gemm.py +++ /dev/null @@ -1,93 +0,0 @@ -from __future__ import annotations - -from contextlib import contextmanager -from typing import TYPE_CHECKING - -import torch - -from sglang.jit_kernel.utils import cache_once, load_jit, override_jit_cuda_arch -from sglang.kernel_api_logging import debug_kernel_api -from sglang.srt.utils.common import is_sm120_supported -from sglang.srt.utils.custom_op import register_custom_op - -if TYPE_CHECKING: - from tvm_ffi.module import Module - - -def _fp8_blockwise_cuda_flags() -> list[str]: - return [ - "-DNDEBUG", - "-DCUTE_USE_PACKED_TUPLE=1", - "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1", - "-DCUTLASS_VERSIONS_GENERATED", - "-DCUTLASS_TEST_LEVEL=0", - "-DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1", - "-DCUTLASS_DEBUG_TRACE_LEVEL=0", - "--expt-relaxed-constexpr", - "--expt-extended-lambda", - ] - - -@contextmanager -def _fp8_blockwise_arch_env(): - if not is_sm120_supported(): - raise RuntimeError( - "fp8_blockwise_scaled_mm JIT kernel requires SM120 (Blackwell)." - ) - major, minor = torch.cuda.get_device_capability() - # sm_*a target (e.g. sm_120a) required, not plain sm_120. - with override_jit_cuda_arch(major, minor, suffix="a"): - yield - - -@cache_once -def _jit_fp8_blockwise_module() -> Module: - """Compile and cache the SM120 fp8 blockwise GEMM module (handles fp16 + bf16).""" - with _fp8_blockwise_arch_env(): - return load_jit( - "fp8_blockwise_scaled_mm", - cuda_files=["gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh"], - cuda_wrappers=[ - ("fp8_blockwise_scaled_mm", "fp8_blockwise_scaled_mm"), - ], - extra_dependencies=["cutlass"], - extra_cuda_cflags=_fp8_blockwise_cuda_flags(), - ) - - -@register_custom_op( - op_name="fp8_blockwise_scaled_mm", - mutates_args=["out"], -) -def _fp8_blockwise_scaled_mm_custom_op( - out: torch.Tensor, - mat_a: torch.Tensor, - mat_b: torch.Tensor, - scales_a: torch.Tensor, - scales_b: torch.Tensor, -) -> None: - module = _jit_fp8_blockwise_module() - module.fp8_blockwise_scaled_mm(out, mat_a, mat_b, scales_a, scales_b) - - -@debug_kernel_api -def fp8_blockwise_scaled_mm( - mat_a: torch.Tensor, - mat_b: torch.Tensor, - scales_a: torch.Tensor, - scales_b: torch.Tensor, - out_dtype: torch.dtype, -) -> torch.Tensor: - """FP8 e4m3 block-wise scaled matmul on SM120.""" - assert out_dtype in ( - torch.float16, - torch.bfloat16, - ), f"out_dtype must be Half or BFloat16, got {out_dtype}" - - out = torch.empty( - (mat_a.shape[0], mat_b.shape[1]), - dtype=out_dtype, - device=mat_a.device, - ) - _fp8_blockwise_scaled_mm_custom_op(out, mat_a, mat_b, scales_a, scales_b) - return out diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh index c5681b933ad9..bd2d9ab52caa 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh @@ -15,7 +15,6 @@ #pragma once -#include #include #include @@ -239,21 +238,6 @@ inline void RuntimeDeviceCheck(DebugInfo location = {}) { return RuntimeDeviceCheck(::cudaGetLastError(), location); } -inline int getSMVersion(int device_id) { - int sm_major = 0; - int sm_minor = 0; - RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id)); - RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id)); - return sm_major * 10 + sm_minor; -} - -inline auto alloc_workspace_tensor(size_t required_bytes, DLDevice device) -> tvm::ffi::Tensor { - if (required_bytes == 0) return {}; - DLDataType u8 = {kDLUInt, 8, 1}; - int64_t shape[] = {static_cast(required_bytes)}; - return ffi::empty(tvm::ffi::ShapeView(shape, 1), u8, device); -} - /** * \brief Kernel launcher with automatic stream resolution and PDL support. * diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index accdd4a4a747..31cffa226579 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -158,9 +158,8 @@ def use_aiter_triton_gemm_w8a8_tuned_gfx950(n: int, k: int) -> bool: if _is_cuda: - from sgl_kernel import fp8_scaled_mm + from sgl_kernel import fp8_blockwise_scaled_mm, fp8_scaled_mm - from sglang.jit_kernel.fp8_blockwise_gemm import fp8_blockwise_scaled_mm from sglang.srt.utils.patch_torch import register_fake_if_exists @register_fake_if_exists("sgl_kernel::fp8_scaled_mm") @@ -170,6 +169,13 @@ def _fp8_scaled_mm_abstract(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=No N = mat_b.shape[-1] return mat_a.new_empty((M, N), dtype=out_dtype) + @register_fake_if_exists("sgl_kernel::fp8_blockwise_scaled_mm") + def _fp8_blockwise_scaled_mm_abstract(mat_a, mat_b, scales_a, scales_b, out_dtype): + # mat_a: [M, K], mat_b: [K, N] or [N, K] depending on callsite layout; output is [M, N]. + M = mat_a.shape[-2] + N = mat_b.shape[-1] + return mat_a.new_empty((M, N), dtype=out_dtype) + use_triton_w8a8_fp8_kernel = get_bool_env_var("USE_TRITON_W8A8_FP8_KERNEL") @@ -268,6 +274,11 @@ def is_aiter(self) -> bool: FP8_GEMM_RUNNER_BACKEND: Fp8GemmRunnerBackend | None = None +def _check_cutlass_block_fp8_hardware_support() -> bool: + """Return True if CUTLASS block FP8 is supported (Hopper or newer with CUDA 12.0+).""" + return is_sm90_supported() or is_blackwell_supported() + + if is_blackwell_supported() and is_flashinfer_available(): from flashinfer import SfLayout from flashinfer import bmm_fp8 as _raw_flashinfer_bmm_fp8 @@ -530,10 +541,11 @@ def _dispatch_explicit_backend(backend: Fp8GemmRunnerBackend) -> Callable: return flashinfer_deepgemm_w8a8_block_fp8_linear_with_fallback elif backend.is_cutlass(): - if not is_sm120_supported(): + if not _check_cutlass_block_fp8_hardware_support(): raise RuntimeError( - "--fp8-gemm-backend=cutlass is deprecated on this hardware. " - "Please switch to DeepGEMM or FlashInfer TRTLLM on SM90/SM100." + "CUTLASS block FP8 requested via --fp8-gemm-backend=cutlass, " + "but hardware does not support it. CUTLASS block FP8 requires " + "Hopper (SM90+) GPUs with CUDA 12.0+." ) return cutlass_w8a8_block_fp8_linear_with_fallback @@ -567,7 +579,7 @@ def _dispatch_auto_backend() -> Callable: # Priority order for auto selection: # 1. DeepGEMM (if enabled and available) # 2. FlashInfer TRTLLM (if Blackwell GPU and FlashInfer available) - # 3. CUTLASS (if SM120 GPU and CUDA 12.8+) + # 3. CUTLASS (if Hopper+ GPU and CUDA 12.0+) # 4. AITER (if AMD GPU with AITER enabled) # 5. Triton (fallback) @@ -575,7 +587,7 @@ def _dispatch_auto_backend() -> Callable: return deepgemm_w8a8_block_fp8_linear_with_fallback elif is_blackwell_supported() and is_flashinfer_available(): return flashinfer_gemm_w8a8_block_fp8_linear_with_fallback - elif is_sm120_supported(): + elif _check_cutlass_block_fp8_hardware_support(): return cutlass_w8a8_block_fp8_linear_with_fallback elif _use_aiter: return aiter_w8a8_block_fp8_linear @@ -589,7 +601,8 @@ def initialize_fp8_gemm_config(server_args: ServerArgs) -> None: backend = server_args.fp8_gemm_runner_backend if backend == "auto" and is_sm120_supported(): - backend = "cutlass" + # TODO(brayden): Verify if CUTLASS can be set by default once SwapAB is supported + backend = "triton" backend = Fp8GemmRunnerBackend(backend) diff --git a/scripts/ci/cuda/warmup_deep_gemm.py b/scripts/ci/cuda/warmup_deep_gemm.py index 270c2e0bd23b..58b7c752c22b 100644 --- a/scripts/ci/cuda/warmup_deep_gemm.py +++ b/scripts/ci/cuda/warmup_deep_gemm.py @@ -115,7 +115,7 @@ def compute_deepseek_v2v3_shapes(config, tp): Shape derivation based on: - MoE: python/sglang/srt/layers/moe/fused_moe_triton/layer.py - MLA: python/sglang/srt/models/deepseek_v2.py - - FP8: python/sglang/kernels/ops/quantization/fp8_kernel.py + - FP8: python/sglang/srt/layers/quantization/fp8_kernel.py """ shapes = [] diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index dc9b4ca27f6e..3c3a9f41641e 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -265,6 +265,7 @@ set(SOURCES "csrc/gemm/awq_kernel.cu" "csrc/gemm/bmm_fp8.cu" "csrc/gemm/dsv3_fused_a_gemm.cu" + "csrc/gemm/fp8_blockwise_gemm_kernel.cu" "csrc/gemm/fp8_gemm_kernel.cu" "csrc/gemm/int8_gemm_kernel.cu" "csrc/gemm/per_token_group_quant_8bit.cu" diff --git a/sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py b/sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py new file mode 100644 index 000000000000..f05687261890 --- /dev/null +++ b/sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py @@ -0,0 +1,237 @@ +import argparse +import copy +import itertools +import os + +import deep_gemm +import torch +import triton +from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor +from sgl_kernel import fp8_blockwise_scaled_mm + +from sglang.utils import is_in_ci + +# Optional vLLM import +try: + from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm + + VLLM_AVAILABLE = True +except ImportError: + vllm_scaled_mm = None + VLLM_AVAILABLE = False + +from sglang.kernels.ops.quantization.fp8_kernel import ( + w8a8_block_fp8_matmul_triton as w8a8_block_fp8_matmul, +) + +IS_CI = is_in_ci() + + +def get_weight_shapes(args): + models_tps = list(itertools.product(args.models, args.tp_sizes)) + # NOTE(HandH1998): The weight shapes only works for DeepSeek-V3. Modify them, if you tune for another different model. + # cannot TP + total = [ + (512 + 64, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (7168, 16384), + (7168, 18432), + ] + # N can TP + n_tp = [ + (18432 * 2, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (24576, 1536), + (4096, 7168), + ] + # K can TP + k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] + # only support Deepseek-V3 + SUPPORT_MODEL = ["deepseek-ai/DeepSeek-V3"] + + weight_shapes = [] + for model, tp_size in models_tps: + assert model in SUPPORT_MODEL + for t in total: + new_t = [t[0], t[1], model] + weight_shapes.append(new_t) + for n_t in n_tp: + new_t = [n_t[0] // tp_size, n_t[1], model] + weight_shapes.append(new_t) + for k_t in k_tp: + new_t = [k_t[0], k_t[1] // tp_size, model] + weight_shapes.append(new_t) + return weight_shapes + + +def cdiv(a: int, b: int) -> int: + """Ceiling division.""" + return -(a // -b) + + +def fp8_gemm_deepgemm( + x_fp8: torch.Tensor, + x_scale: torch.Tensor, + y_fp8: torch.Tensor, + y_scale: torch.Tensor, + m: int, + n: int, + k: int, +): + """DeepGEMM implementation of FP8 GEMM""" + out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + + # Run DeepGEMM kernel + deep_gemm.fp8_gemm_nt((x_fp8, x_scale), (y_fp8, y_scale), out) + return out + + +def scale_shape(shape, group_shape): + assert len(shape) == len(group_shape) + return tuple(cdiv(shape[i], group_shape[i]) for i in range(len(group_shape))) + + +# CI environment uses simplified parameters +if IS_CI: + batch_sizes = [1, 8] # Simplified for CI +else: + batch_sizes = [1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096] + +# Filter providers based on availability +available_providers = ["sgl-kernel"] +available_names = ["sgl-kernel"] +available_styles = [("orange", "-")] + +if VLLM_AVAILABLE: + available_providers.insert(0, "vllm") + available_names.insert(0, "vllm") + available_styles.insert(0, ("blue", "-")) + +available_providers.append("triton") +available_names.append("sglang triton") +available_styles.append(("red", "-")) + +# Add deepgemm if available +try: + import deep_gemm + + available_providers.append("deepgemm") + available_names.append("deepgemm") + available_styles.append(("yellow", "-")) +except ImportError: + pass + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size"], + x_vals=batch_sizes, + x_log=False, + line_arg="provider", + line_vals=available_providers, + line_names=available_names, + styles=available_styles, + ylabel="GB/s", + plot_name="fp8 blockwise scaled matmul", + args={}, + ) +) +def benchmark(batch_size, provider, N, K): + M = batch_size + fp8_info = torch.finfo(torch.float8_e4m3fn) + fp8_max, fp8_min = fp8_info.max, fp8_info.min + + a_fp32 = (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max + a_fp8 = a_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) + + b_fp32 = (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max + b_fp8 = b_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) + + scale_a_group_shape = (1, 128) + scale_b_group_shape = (128, 128) + scale_a_shape = scale_shape(a_fp8.shape, scale_a_group_shape) + scale_b_shape = scale_shape(b_fp8.shape, scale_b_group_shape) + + scale_a = torch.randn(scale_a_shape, device="cuda", dtype=torch.float32) + scale_b = torch.randn(scale_b_shape, device="cuda", dtype=torch.float32) + + quantiles = [0.5, 0.2, 0.8] + if provider == "sgl-kernel": + scale_a = scale_a.t().contiguous().t() + b_fp8, scale_b = b_fp8.t(), scale_b.t() + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + lambda: fp8_blockwise_scaled_mm( + a_fp8, b_fp8, scale_a, scale_b, torch.float16 + ), + quantiles=quantiles, + ) + elif provider == "vllm": + if not VLLM_AVAILABLE: + return (0, 0, 0) + scale_a = scale_a.t().contiguous().t() + b_fp8, scale_b = b_fp8.t(), scale_b.t() + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + lambda: vllm_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, torch.float16), + quantiles=quantiles, + ) + elif provider == "triton": + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + lambda: w8a8_block_fp8_matmul( + a_fp8, b_fp8, scale_a, scale_b, [128, 128], torch.float16 + ), + quantiles=quantiles, + ) + if provider == "deepgemm": + scale_a_col_major = get_mn_major_tma_aligned_tensor(scale_a.clone()) + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + lambda: fp8_gemm_deepgemm( + a_fp8, scale_a_col_major, b_fp8, scale_b, M, N, K + ), + quantiles=quantiles, + ) + return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--models", + nargs="+", + type=str, + default=["deepseek-ai/DeepSeek-V3"], + help="List of models to benchmark", + ) + parser.add_argument( + "--tp-sizes", + nargs="+", + type=int, + default=[1], + help="List of tensor parallel sizes", + ) + args = parser.parse_args() + + # Simplify for CI environment + if IS_CI: + args.models = [args.models[0]] # Use only first model + args.tp_sizes = [args.tp_sizes[0]] # Use only first TP size + + NK_model_names = get_weight_shapes(args) + + # Limit iterations in CI + if IS_CI: + NK_model_names = NK_model_names[:2] # Only test first 2 shapes in CI + + for N, K, model_name in NK_model_names: + if N % 128 != 0 or K % 128 != 0: + print(f"Skip {N=}, {K=} now") + continue + print(f"{model_name} N={N} K={K}: ") + benchmark.run( + print_data=True, + N=N, + K=K, + ) + + print("Benchmark finished!") diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index 134aaf453b1d..45d3dfe27b55 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -123,6 +123,11 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "bias) -> Tensor"); m.impl("fp8_scaled_mm", torch::kCUDA, &fp8_scaled_mm); + m.def( + "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> " + "Tensor"); + m.impl("fp8_blockwise_scaled_mm", torch::kCUDA, &fp8_blockwise_scaled_mm); + m.def( "sgl_per_token_group_quant_8bit(Tensor input, Tensor! output_q, Tensor! output_s, int group_size," " float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()"); diff --git a/sgl-kernel/csrc/cutlass_extensions/gemm/fp8_blockwise_gemm_sm90_dispatch.cuh b/sgl-kernel/csrc/cutlass_extensions/gemm/fp8_blockwise_gemm_sm90_dispatch.cuh new file mode 100644 index 000000000000..05b70c4f26f2 --- /dev/null +++ b/sgl-kernel/csrc/cutlass_extensions/gemm/fp8_blockwise_gemm_sm90_dispatch.cuh @@ -0,0 +1,197 @@ +// Adapted from +// https://github.com/vllm-project/vllm/blob/main/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/kernel/tile_scheduler_params.h" +#include "cutlass/numeric_types.h" +#include "cutlass/tensor_ref.h" +#include "cutlass_extensions/common.hpp" +#include "cutlass_extensions/gemm/cutlass_gemm_caller.cuh" +#include "cutlass_extensions/gemm/dispatch_policy.hpp" + +using namespace cute; + +template < + typename SchedulerType, + typename OutType, + int GroupSizeM_, + int GroupSizeN_, + int GroupSizeK_, + int TileSizeM_ = 128, + class ClusterShape = Shape<_1, _2, _1>> +struct cutlass_3x_gemm_fp8_blockwise { + using GroupSizeM = Int; + using GroupSizeN = Int; + using GroupSizeK = Int; + using TileSizeM = Int; + + static_assert(TileSizeM_ % GroupSizeM_ == 0, "TileSizeM must be a multiple of GroupSizeM"); + + using ElementAB = cutlass::float_e4m3_t; + + // A matrix configuration + using ElementA = ElementAB; + using LayoutA = cutlass::layout::RowMajor; + static constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; + + // B matrix configuration + using ElementB = ElementAB; + using LayoutB = cutlass::layout::ColumnMajor; + static constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; + + // C/D matrix configuration + using ElementC = void; + using LayoutC = cutlass::layout::RowMajor; + static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + + using ElementD = OutType; + using LayoutD = cutlass::layout::RowMajor; + static constexpr int AlignmentD = AlignmentC; + + using ScaleTileShape = Shape<_1, _128, _128>; + using ScaleConfig = decltype(cutlass::detail::sm90_trivial_blockwise_scale_config(ScaleTileShape{})); + using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); + using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); + + // Multiply-accumulate blocking/pipelining details + using ElementAccumulator = float; // Element type for internal accumulation + using ElementCompute = float; // Element type for compute + using TileShape = Shape; // Threadblock-level tile size + + using ArchTag = cutlass::arch::Sm90; + using OperatorClass = cutlass::arch::OpClassTensorOp; + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative; + using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; + using StoreEpilogueCompute = typename cutlass::epilogue::fusion::Sm90EVT; + + using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8Blockwise; + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + TileShape, + ClusterShape, + EpilogueTileType, + ElementAccumulator, + ElementCompute, + ElementC, + LayoutC, + AlignmentC, + ElementD, + LayoutD, + AlignmentD, + EpilogueSchedule, + StoreEpilogueCompute>::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + ElementA, + cute::tuple, + AlignmentA, + ElementB, + cute::tuple, + AlignmentB, + ElementAccumulator, + TileShape, + ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelSchedule>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, // Indicates ProblemShape + CollectiveMainloop, + CollectiveEpilogue, + SchedulerType>; +}; + +template +void cutlass_gemm_caller_blockwise( + torch::Tensor& out, + torch::Tensor const& a, + torch::Tensor const& b, + torch::Tensor const& a_scales, + torch::Tensor const& b_scales) { + using GemmKernel = typename Gemm::GemmKernel; + using ElementAB = typename Gemm::ElementAB; + using ElementA = ElementAB; + using ElementB = ElementAB; + using ElementD = typename Gemm::ElementD; + using ElementBlockScale = float; + + using ScaleTileShape = Shape<_1, _128, _128>; + using ScaleConfig = decltype(cutlass::detail::sm90_trivial_blockwise_scale_config(ScaleTileShape{})); + using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); + using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); + + int m = a.size(0); + int k = a.size(1); + int n = b.size(1); + + auto a_ptr = static_cast(a.data_ptr()); + auto b_ptr = static_cast(b.data_ptr()); + + auto a_s_ptr = static_cast(a_scales.data_ptr()); + auto b_s_ptr = static_cast(b_scales.data_ptr()); + + using StrideA = typename GemmKernel::StrideA; + using StrideB = typename GemmKernel::StrideB; + using StrideD = typename GemmKernel::StrideD; + using StrideC = typename GemmKernel::StrideC; + + StrideA a_stride = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1)); + StrideB b_stride = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1)); + StrideC c_stride = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1)); + LayoutSFA layout_sfa = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1)); + LayoutSFB layout_sfb = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1)); + + typename GemmKernel::MainloopArguments mainloop_args{ + a_ptr, a_stride, b_ptr, b_stride, a_s_ptr, layout_sfa, b_s_ptr, layout_sfb}; + auto c_ptr = static_cast(out.data_ptr()); + typename GemmKernel::EpilogueArguments epilogue_args{{}, c_ptr, c_stride, c_ptr, c_stride}; + + typename GemmKernel::TileSchedulerArguments scheduler; + + static constexpr bool UsesStreamKScheduler = + cute::is_same_v; + + if constexpr (UsesStreamKScheduler) { + using DecompositionMode = + typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90StreamKParams::DecompositionMode; + using ReductionMode = + typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90StreamKParams::ReductionMode; + + scheduler.decomposition_mode = DecompositionMode::StreamK; + scheduler.reduction_mode = ReductionMode::Nondeterministic; + } + + cutlass_gemm_caller(a.device(), {m, n, k, 1}, mainloop_args, epilogue_args, scheduler); +} + +template +void cutlass_gemm_blockwise_sm90_fp8_dispatch( + torch::Tensor& out, + torch::Tensor const& a, + torch::Tensor const& b, + torch::Tensor const& a_scales, + torch::Tensor const& b_scales) { + auto k = a.size(1); + auto n = b.size(1); + + if (k > 3 * n) { + cutlass_gemm_caller_blockwise>( + out, a, b, a_scales, b_scales); + } else { + cutlass_gemm_caller_blockwise< + cutlass_3x_gemm_fp8_blockwise>( + out, a, b, a_scales, b_scales); + } +} diff --git a/sgl-kernel/csrc/gemm/fp8_blockwise_gemm_kernel.cu b/sgl-kernel/csrc/gemm/fp8_blockwise_gemm_kernel.cu new file mode 100644 index 000000000000..cc094de51a60 --- /dev/null +++ b/sgl-kernel/csrc/gemm/fp8_blockwise_gemm_kernel.cu @@ -0,0 +1,522 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cutlass_extensions/gemm/cutlass_gemm_caller.cuh" +#include "cutlass_extensions/gemm/fp8_blockwise_gemm_sm90_dispatch.cuh" +#include "utils.h" + +using namespace cute; + +template < + typename OutType, + typename MmaTileShape, + typename PerSmTileShape, + typename EpilogueTileShape, + typename ScalesPerTile, + int TileSizeM_ = 128, + class ClusterShape = Shape<_1, _1, _1>> +void launch_sm100_fp8_blockwise_scaled_mm( + torch::Tensor& out, + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& scales_a, + const torch::Tensor& scales_b) { + static constexpr int ScaleMsPerTile = size<0>(ScalesPerTile{}); + static constexpr int ScaleGranularityM = size<0>(MmaTileShape{}) / ScaleMsPerTile; + static constexpr int ScaleGranularityN = size<1>(MmaTileShape{}) / size<1>(ScalesPerTile{}); + static constexpr int ScaleGranularityK = size<2>(MmaTileShape{}) / size<2>(ScalesPerTile{}); + + using ElementAB = cutlass::float_e4m3_t; + using ElementA = ElementAB; + using ElementB = ElementAB; + using ElementC = void; + using ElementD = OutType; + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutD = cutlass::layout::RowMajor; + using LayoutC = LayoutD; + // This means both SFA and SFB are column-major. + using ScaleConfig = cutlass::detail::Sm100BlockwiseScaleConfig< + ScaleGranularityM, + ScaleGranularityN, + ScaleGranularityK, + cute::UMMA::Major::MN, + cute::UMMA::Major::K>; + using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); + using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); + + static constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; + static constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; + static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + static constexpr int AlignmentC = AlignmentD; + + using ElementAccumulator = float; + using ElementBlockScale = float; + using ElementCompute = float; + using ArchTag = cutlass::arch::Sm100; + using OperatorClass = cutlass::arch::OpClassTensorOp; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, + cutlass::arch::OpClassTensorOp, + PerSmTileShape, + ClusterShape, + EpilogueTileShape, + ElementAccumulator, + ElementCompute, + ElementC, + LayoutC, + AlignmentC, + ElementD, + LayoutD, + AlignmentD, + cutlass::epilogue::TmaWarpSpecialized1Sm>::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + ElementA, + cute::tuple, + AlignmentA, + ElementB, + cute::tuple, + AlignmentB, + ElementAccumulator, + MmaTileShape, + ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecializedBlockwise1SmSm100>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::PersistentScheduler>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + Gemm gemm_op; + + int m = a.size(0); + int k = a.size(1); + int n = b.size(1); + + auto a_ptr = static_cast(a.data_ptr()); + auto b_ptr = static_cast(b.data_ptr()); + auto scales_a_ptr = static_cast(scales_a.data_ptr()); + auto scales_b_ptr = static_cast(scales_b.data_ptr()); + auto c_ptr = static_cast(out.data_ptr()); + + using StrideA = typename GemmKernel::StrideA; + using StrideB = typename GemmKernel::StrideB; + using StrideD = typename GemmKernel::StrideD; + using StrideC = typename GemmKernel::StrideD; + + StrideA a_stride = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1)); + StrideB b_stride = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1)); + StrideC c_stride = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1)); + LayoutSFA layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1)); + LayoutSFB layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1)); + + typename GemmKernel::MainloopArguments mainloop_args{ + a_ptr, a_stride, b_ptr, b_stride, scales_a_ptr, layout_SFA, scales_b_ptr, layout_SFB}; + + typename GemmKernel::EpilogueArguments epilogue_args{{}, c_ptr, c_stride, c_ptr, c_stride}; + epilogue_args.thread.alpha = 1.0f; + + typename GemmKernel::Arguments args = { + cutlass::gemm::GemmUniversalMode::kGemm, {m, n, k, 1}, mainloop_args, epilogue_args}; + + auto can_implement = gemm_op.can_implement(args); + TORCH_CHECK(can_implement == cutlass::Status::kSuccess, cutlassGetStatusString(can_implement)) + + size_t workspace_size = gemm_op.get_workspace_size(args); + cutlass::device_memory::allocation workspace(workspace_size); + + auto init_status = gemm_op.initialize(args, workspace.get()); + TORCH_CHECK(init_status == cutlass::Status::kSuccess, cutlassGetStatusString(init_status)); + + auto stream = at::cuda::getCurrentCUDAStream(a.get_device()); + + auto status = gemm_op.run(stream); + TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)) +} + +template +void sm100_fp8_blockwise_dispatch_shape( + torch::Tensor& out, + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& scales_a, + const torch::Tensor& scales_b) { + if (a.size(0) <= 128) { + using MmaTileShape = Shape<_64, _128, _128>; + using PerSmTileShape = Shape<_64, _128, _128>; + using EpilogueTileShape = Shape<_64, _64>; + using ScalesPerTile = Shape<_64, _1, _1>; + launch_sm100_fp8_blockwise_scaled_mm( + out, a, b, scales_a, scales_b); + } else { + using MmaTileShape = Shape<_128, _128, _128>; + using PerSmTileShape = Shape<_128, _128, _128>; + using EpilogueTileShape = Shape<_128, _64>; + using ScalesPerTile = Shape<_128, _1, _1>; + launch_sm100_fp8_blockwise_scaled_mm( + out, a, b, scales_a, scales_b); + } +} + +template < + typename OutType, + typename MmaTileShape, + typename PerSmTileShape, + typename EpilogueTileShape, + typename ScalesPerTile, + int TileSizeM_ = 128, + class ClusterShape = Shape<_1, _1, _1>> +void launch_sm120_fp8_blockwise_scaled_mm( + torch::Tensor& out, + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& scales_a, + const torch::Tensor& scales_b) { + using ElementBlockScale = float; + + // A matrix configuration + using ElementA = cutlass::float_e4m3_t; // Element type for A matrix operand + using LayoutATag = cutlass::layout::RowMajor; // Layout type for A matrix operand + constexpr int AlignmentA = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of A matrix in units of + // elements (up to 16 bytes) + + // B matrix configuration + using ElementB = cutlass::float_e4m3_t; // Element type for B matrix operand + using LayoutBTag = cutlass::layout::ColumnMajor; // Layout type for B matrix operand + constexpr int AlignmentB = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of B matrix in units of + // elements (up to 16 bytes) + + // C/D matrix configuration + using ElementD = OutType; // Element type for D matrix operand + using ElementC = void; // Element type for C matrix operand + using LayoutCTag = cutlass::layout::RowMajor; // Layout type for C matrix operand + using LayoutDTag = cutlass::layout::RowMajor; // Layout type for D matrix operand + constexpr int AlignmentD = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of C matrix in units of + // elements (up to 16 bytes) + constexpr int AlignmentC = + AlignmentD; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes) + + // Kernel functional config + using ElementAccumulator = float; // Element type for internal accumulation + using ArchTag = cutlass::arch::Sm120; // Tag indicating the minimum SM that supports the intended feature + using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag - changed from OpClassBlockScaledTensorOp + + static constexpr int ScaleMsPerTile = size<0>(ScalesPerTile{}); + static constexpr int ScaleGranularityM = size<0>(MmaTileShape{}) / ScaleMsPerTile; + static constexpr int ScaleGranularityN = size<1>(MmaTileShape{}) / size<1>(ScalesPerTile{}); + static constexpr int ScaleGranularityK = size<2>(MmaTileShape{}) / size<2>(ScalesPerTile{}); + + using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig< + ScaleGranularityM, + ScaleGranularityN, + ScaleGranularityK, + cute::UMMA::Major::MN, + cute::UMMA::Major::K>; + // FP8 Block-wise scaling configuration + using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); // Layout type for SFA matrix operand + using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); // Layout type for SFB matrix operand + + constexpr bool kCanUsePingpong = (64 % ScaleGranularityM == 0); + + int m = a.size(0); + int k = a.size(1); + int n = b.size(1); + + auto a_ptr = static_cast(a.data_ptr()); + auto b_ptr = static_cast(b.data_ptr()); + auto c_ptr = static_cast(out.data_ptr()); + + auto scales_a_ptr = static_cast(scales_a.data_ptr()); + auto scales_b_ptr = static_cast(scales_b.data_ptr()); + + LayoutSFA layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1)); + LayoutSFB layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1)); + + auto run_gemm = [&](auto tag) -> cutlass::Status { + using GemmKernel = decltype(tag); + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + Gemm gemm_op; + + using StrideA = typename GemmKernel::StrideA; + using StrideB = typename GemmKernel::StrideB; + using StrideC = typename GemmKernel::StrideD; + + StrideA stride_a = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1)); + StrideB stride_b = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1)); + StrideC stride_c = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1)); + + typename GemmKernel::MainloopArguments mainloop_args{ + a_ptr, stride_a, b_ptr, stride_b, scales_a_ptr, layout_SFA, scales_b_ptr, layout_SFB}; + + typename GemmKernel::EpilogueArguments epilogue_args{{}, c_ptr, stride_c, c_ptr, stride_c}; + epilogue_args.thread.alpha = 1.0f; + + typename Gemm::Arguments args = { + cutlass::gemm::GemmUniversalMode::kGemm, + {m, n, k, 1}, + mainloop_args, + epilogue_args, + }; + + auto can_implement = gemm_op.can_implement(args); + if (can_implement != cutlass::Status::kSuccess) { + return can_implement; + } + + size_t workspace_size = gemm_op.get_workspace_size(args); + cutlass::device_memory::allocation workspace(workspace_size); + + auto init_status = gemm_op.initialize(args, workspace.get()); + if (init_status != cutlass::Status::kSuccess) { + return init_status; + } + + auto stream = at::cuda::getCurrentCUDAStream(a.get_device()); + return gemm_op.run(stream); + }; + + using CooperativeCollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + PerSmTileShape, + ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, + ElementAccumulator, + ElementC, + LayoutCTag, + AlignmentC, + ElementD, + LayoutDTag, + AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; + + using CooperativeStageCount = cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CooperativeCollectiveEpilogue::SharedStorage))>; + + using CooperativeCollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + ElementA, + cute::tuple, + AlignmentA, + ElementB, + cute::tuple, + AlignmentB, + ElementAccumulator, + MmaTileShape, + ClusterShape, + CooperativeStageCount, + cutlass::gemm::KernelScheduleSm120Blockwise>::CollectiveOp; + + using CooperativeGemmKernel = cutlass::gemm::kernel:: + GemmUniversal, CooperativeCollectiveMainloop, CooperativeCollectiveEpilogue, void>; + + cutlass::Status status = cutlass::Status::kSuccess; + if constexpr (kCanUsePingpong) { + using PingpongMmaTileShape_MNK = Shape<_64, _128, _128>; + using PingpongCollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + PerSmTileShape, + ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, + ElementAccumulator, + ElementC, + LayoutCTag, + AlignmentC, + ElementD, + LayoutDTag, + AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; + + using PingpongStageCount = cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename PingpongCollectiveEpilogue::SharedStorage))>; + + using PingpongCollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + ElementA, + cute::tuple, + AlignmentA, + ElementB, + cute::tuple, + AlignmentB, + ElementAccumulator, + PingpongMmaTileShape_MNK, + ClusterShape, + PingpongStageCount, + cutlass::gemm::KernelTmaWarpSpecializedBlockwisePingpongSm120>::CollectiveOp; + + using PingpongGemmKernel = cutlass::gemm::kernel:: + GemmUniversal, PingpongCollectiveMainloop, PingpongCollectiveEpilogue, void>; + + if (m <= 64) { + status = run_gemm(PingpongGemmKernel{}); + if (status != cutlass::Status::kSuccess) { + status = run_gemm(CooperativeGemmKernel{}); + } + } else { + status = run_gemm(CooperativeGemmKernel{}); + } + } else { + status = run_gemm(CooperativeGemmKernel{}); + } + + TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); +} + +template +void sm120_fp8_blockwise_dispatch_shape( + torch::Tensor& out, + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& scales_a, + const torch::Tensor& scales_b) { + using MmaTileShape = Shape<_128, _128, _128>; + using PerSmTileShape = Shape<_128, _128, _128>; + using EpilogueTileShape = Shape<_128, _64>; + using ScalesPerTile = Shape<_128, _1, _1>; + launch_sm120_fp8_blockwise_scaled_mm( + out, a, b, scales_a, scales_b); +} + +torch::Tensor fp8_blockwise_scaled_mm( + const torch::Tensor& mat_a, + const torch::Tensor& mat_b, + const torch::Tensor& scales_a, + const torch::Tensor& scales_b, + const torch::Dtype& out_dtype) { + TORCH_CHECK(mat_a.is_cuda(), "mat_a must be a CUDA tensor"); + TORCH_CHECK(mat_b.is_cuda(), "mat_b must be a CUDA tensor"); + TORCH_CHECK(mat_a.dim() == 2, "mat_a must be a 2D tensor"); + TORCH_CHECK(mat_b.dim() == 2, "mat_b must be a 2D tensor"); + TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor"); + TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor"); + TORCH_CHECK(mat_a.size(1) == mat_b.size(0), "mat_a and mat_b shapes cannot be multiplied"); + + TORCH_CHECK( + (mat_a.size(1) * mat_a.element_size()) % 16 == 0, "mat_a must be multiple of 16 bytes for memory alignment"); + TORCH_CHECK( + (mat_b.size(0) * mat_b.element_size()) % 16 == 0, "mat_b must be multiple of 16 bytes for memory alignment"); + TORCH_CHECK(mat_a.scalar_type() == torch::kFloat8_e4m3fn, "mat_a must be Float8_e4m3fn"); + TORCH_CHECK(mat_b.scalar_type() == torch::kFloat8_e4m3fn, "mat_b must be Float8_e4m3fn"); + TORCH_CHECK(out_dtype == torch::kHalf || out_dtype == torch::kBFloat16, "out_dtype must be Half or BFloat16"); + + auto is_contiguous_vector = [](const torch::Tensor& t) { + auto t_sizes = t.sizes(); + return t.is_contiguous() && + (t.dim() == 1 || (t.dim() == 2 && *std::min_element(t_sizes.begin(), t_sizes.end()) == 1)); + }; + + TORCH_CHECK(mat_a.size(0) == scales_a.size(0), "size of scales_a is not matched"); + TORCH_CHECK(mat_a.size(1) / 128 == scales_a.size(1), "size of scales_a is not matched"); + TORCH_CHECK(scales_a.stride(0) == 1 || is_contiguous_vector(scales_a), "scales_a must be M major"); + TORCH_CHECK(mat_b.size(0) / 128 == scales_b.size(0), "size of scales_b is not matched"); + TORCH_CHECK(mat_b.size(1) / 128 == scales_b.size(1), "size of scales_b is not matched"); + TORCH_CHECK(scales_b.stride(0) == 1 || is_contiguous_vector(scales_b), "scales_b must be K major"); + TORCH_CHECK(scales_a.scalar_type() == torch::kFloat32, "scales_a must be Float32"); + TORCH_CHECK(scales_b.scalar_type() == torch::kFloat32, "scales_b must be Float32"); + + torch::Tensor out = torch::empty({mat_a.size(0), mat_b.size(1)}, mat_a.options().dtype(out_dtype)); + TORCH_CHECK((out.size(1) * out.element_size()) % 16 == 0, "out must be multiple of 16 bytes for memory alignment"); + + auto sm_version = getSMVersion(); + + int64_t original_rows = mat_a.size(0); + torch::Tensor mat_a_padded = pad_tensor(mat_a, /*alignment=*/4); + torch::Tensor scales_a_padded = pad_tensor(scales_a, /*alignment=*/4, /*col_major=*/true); + torch::Tensor out_padded = torch::empty({mat_a_padded.size(0), mat_b.size(1)}, out.options()); + +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) +#if defined CUDA_VERSION && CUDA_VERSION >= 12000 + if (sm_version == 90) { + torch::Tensor scales_b_contiguous = scales_b.contiguous(); + if (out_dtype == torch::kBFloat16) { + cutlass_gemm_blockwise_sm90_fp8_dispatch( + out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b_contiguous); + } else { + cutlass_gemm_blockwise_sm90_fp8_dispatch( + out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b_contiguous); + } + return out_padded.slice(0, 0, original_rows); + } +#endif +#endif + +#if defined(CUTLASS_ARCH_MMA_SM100A_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) +#if defined CUDA_VERSION && CUDA_VERSION >= 12080 + if (sm_version == 100 +#if CUDA_VERSION >= 12090 + || sm_version == 103 +#endif + ) { + if (out_dtype == torch::kBFloat16) { + sm100_fp8_blockwise_dispatch_shape( + out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b); + } else { + sm100_fp8_blockwise_dispatch_shape(out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b); + } + return out_padded.slice(0, 0, original_rows); + } +#endif +#endif + +#if defined(CUTLASS_ARCH_MMA_SM120A_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) +#if defined(CUDA_VERSION) && CUDA_VERSION >= 12080 + if (sm_version >= 120) { + if (out_dtype == torch::kBFloat16) { + sm120_fp8_blockwise_dispatch_shape( + out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b); + } else { + sm120_fp8_blockwise_dispatch_shape(out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b); + } + return out_padded.slice(0, 0, original_rows); + } +#endif +#endif + + TORCH_CHECK_NOT_IMPLEMENTED( + false, "No implemented fp8_blockwise_scaled_mm for current compute capability: ", sm_version); +} diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 3f9cda5029ee..9a92f38bf6bb 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -235,6 +235,12 @@ torch::Tensor fp8_scaled_mm( const torch::Tensor& scales_b, const torch::Dtype& out_dtype, const c10::optional& bias); +torch::Tensor fp8_blockwise_scaled_mm( + const torch::Tensor& mat_a, + const torch::Tensor& mat_b, + const torch::Tensor& scales_a, + const torch::Tensor& scales_b, + const torch::Dtype& out_dtype); void sgl_per_token_group_quant_8bit( at::Tensor input, at::Tensor output_q, diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index 21f5eb90f487..199c1f92264d 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -57,6 +57,7 @@ awq_dequantize, bmm_fp8, dsv3_fused_a_gemm, + fp8_blockwise_scaled_mm, fp8_scaled_mm, gptq_gemm, gptq_shuffle, @@ -177,6 +178,7 @@ "fast_topk_transform_ragged_fused", "fast_topk_v2", "fp8_blockwise_scaled_grouped_mm", + "fp8_blockwise_scaled_mm", "fp8_scaled_mm", "fused_add_rmsnorm", "fused_qk_norm_rope", diff --git a/sgl-kernel/python/sgl_kernel/gemm.py b/sgl-kernel/python/sgl_kernel/gemm.py index 56d88ae14d99..1d68cdf94943 100644 --- a/sgl-kernel/python/sgl_kernel/gemm.py +++ b/sgl-kernel/python/sgl_kernel/gemm.py @@ -21,6 +21,16 @@ def int8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None): ) +def fp8_blockwise_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype): + return torch.ops.sgl_kernel.fp8_blockwise_scaled_mm.default( + mat_a, + mat_b, + scales_a, + scales_b, + out_dtype, + ) + + def fp8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None): return torch.ops.sgl_kernel.fp8_scaled_mm.default( mat_a, diff --git a/test/registered/jit/test_fp8_blockwise_gemm.py b/sgl-kernel/tests/test_fp8_blockwise_gemm.py similarity index 69% rename from test/registered/jit/test_fp8_blockwise_gemm.py rename to sgl-kernel/tests/test_fp8_blockwise_gemm.py index 5d097b9d28d3..a4438de4afc4 100644 --- a/test/registered/jit/test_fp8_blockwise_gemm.py +++ b/sgl-kernel/tests/test_fp8_blockwise_gemm.py @@ -1,18 +1,11 @@ +import os +import random import sys from typing import Optional, Type import pytest import torch - -from sglang.jit_kernel.fp8_blockwise_gemm import fp8_blockwise_scaled_mm -from sglang.srt.utils import is_sm120_supported -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci( - est_time=30, - stage="base-b", - runner_config="1-gpu-small", -) +from sgl_kernel import fp8_blockwise_scaled_mm def cdiv(a: int, b: int) -> int: @@ -32,6 +25,20 @@ def baseline_scaled_mm( out_dtype: Type[torch.dtype], bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: + # We treat N-dimensional group scaling as extended numpy-style broadcasting + # in numpy simply stretches dimensions with an extent of 1 to match the + # the target shape by repeating the data along that dimension (broadcasting) + # , we extend these semantics to say if the extent of a dimension in the + # source shape is not 1 and does not match the target shape we repeat each + # element along that dimension src_shape[dim] // target_shape[dim] times + # example if we have: + # a = [[1, 2], and target_shape = (2, 4) + # [3, 4]] + # then we would expand a to: + # a = [[1, 1, 2, 2], + # [3, 3, 4, 4]] + # NOTE this function this function does not explicitly broadcast dimensions + # with an extent of 1, since this can be done implicitly by pytorch def group_broadcast(t, shape): for i, s in enumerate(shape): if t.shape[i] != s and t.shape[i] != 1: @@ -75,16 +82,13 @@ def _test_accuracy_once(M, N, K, out_dtype, device): torch.testing.assert_close(o, o1, rtol=rtol, atol=atol) -@pytest.mark.skipif( - not is_sm120_supported(), reason="fp8_blockwise_scaled_mm requires SM120 (>= 12.0)" -) -@pytest.mark.parametrize("M", [1, 3, 5, 32, 48, 64, 127, 128, 512, 1024, 4096]) -@pytest.mark.parametrize("N", [128, 512, 1024, 4096, 8192]) -@pytest.mark.parametrize("K", [512, 1024, 4096, 8192]) +@pytest.mark.parametrize("M", [1, 3, 5, 127, 128, 512, 1024, 4096]) +@pytest.mark.parametrize("N", [128, 512, 1024, 4096, 8192, 14080]) +@pytest.mark.parametrize("K", [512, 1024, 4096, 8192, 14080, 16384]) @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16]) def test_accuracy(M, N, K, out_dtype): _test_accuracy_once(M, N, K, out_dtype, "cuda") if __name__ == "__main__": - sys.exit(pytest.main([__file__, "-v"])) + sys.exit(pytest.main([__file__])) diff --git a/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py b/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py deleted file mode 100644 index dcab974f154a..000000000000 --- a/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py +++ /dev/null @@ -1,103 +0,0 @@ -from __future__ import annotations - -import sys - -import torch -import triton - -from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark -from sglang.jit_kernel.fp8_blockwise_gemm import fp8_blockwise_scaled_mm -from sglang.srt.utils import is_sm120_supported -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci( - est_time=5, - stage="base-b-kernel-benchmark", - runner_config="1-gpu-large", -) - - -def _make_inputs(m: int, n: int, k: int, device: str = "cuda"): - fp8_info = torch.finfo(torch.float8_e4m3fn) - fp8_max, fp8_min = fp8_info.max, fp8_info.min - a_fp32 = (torch.rand(m, k, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max - a_fp8 = a_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) - b_fp32 = (torch.rand(n, k, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max - b_fp8 = b_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn).t() - - scale_a = torch.randn((m, k // 128), device=device, dtype=torch.float32) * 0.001 - scale_b = ( - torch.randn((k // 128, n // 128), device=device, dtype=torch.float32) * 0.001 - ) - scale_a = scale_a.t().contiguous().t() - scale_b = scale_b.t().contiguous().t() - return a_fp8, b_fp8, scale_a, scale_b - - -def _torch_ref(a_fp8, b_fp8, scale_a, scale_b): - def group_broadcast(t, shape): - for i, s in enumerate(shape): - if t.shape[i] != s and t.shape[i] != 1: - assert s % t.shape[i] == 0 - t = ( - t.unsqueeze(i + 1) - .expand(*t.shape[: i + 1], s // t.shape[i], *t.shape[i + 1 :]) - .flatten(i, i + 1) - ) - return t - - sa = group_broadcast(scale_a, a_fp8.shape) - sb = group_broadcast(scale_b, b_fp8.shape) - return torch.mm(sa * a_fp8.to(torch.float32), sb * b_fp8.to(torch.float32)).to( - torch.bfloat16 - ) - - -shape_range = get_benchmark_range( - full_range=[ - (16, 4096, 4096), # swapAB tile N=32 - (64, 4096, 4096), # swapAB tile N=64 - (128, 4096, 4096), # non-swap 128 - (512, 4096, 4096), - (1024, 8192, 4096), - ], - ci_range=[(16, 4096, 4096), (128, 4096, 4096)], -) - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["m", "n", "k"], - x_vals=shape_range, - x_log=False, - line_arg="provider", - line_vals=["jit", "torch_ref"], - line_names=["JIT FP8 Blockwise GEMM", "Torch Ref"], - styles=[("green", "-"), ("blue", "-")], - ylabel="us", - plot_name="fp8-blockwise-scaled-mm-performance", - args={}, - ) -) -def benchmark(m, n, k, provider): - a_fp8, b_fp8, scale_a, scale_b = _make_inputs(m, n, k) - - if provider == "jit": - fn = lambda: fp8_blockwise_scaled_mm( - a_fp8, b_fp8, scale_a, scale_b, out_dtype=torch.bfloat16 - ) - elif provider == "torch_ref": - fn = lambda: _torch_ref(a_fp8, b_fp8, scale_a, scale_b) - else: - raise ValueError(f"Unknown provider: {provider}") - - return run_benchmark(fn) - - -if __name__ == "__main__": - if not is_sm120_supported(): - print( - "[skip] fp8_blockwise_scaled_mm benchmark requires SM120 with CUDA 12.8+." - ) - sys.exit(0) - benchmark.run(print_data=True) diff --git a/test/registered/quant/test_fp8_blockwise_row_padding.py b/test/registered/quant/test_fp8_blockwise_row_padding.py new file mode 100644 index 000000000000..43b49b5ecc09 --- /dev/null +++ b/test/registered/quant/test_fp8_blockwise_row_padding.py @@ -0,0 +1,136 @@ +"""Unit tests for the row-padded quant path of the cutlass FP8 blockwise linear. + +`cutlass_w8a8_block_fp8_linear_with_fallback` quantizes activations into +row-aligned buffers (`sglang_per_token_group_quant_fp8_row_padded`) so the +`fp8_blockwise_scaled_mm` wrapper's per-call mat_a/scales_a padding short- +circuits. These tests pin the invariant that this is numerically identical to +the legacy unpadded path, across both row-aligned and unaligned M. +""" + +import unittest + +import torch + +from sglang.kernels.ops.quantization.fp8_kernel import ( + fp8_dtype, + per_token_group_quant_fp8, + sglang_per_token_group_quant_fp8_row_padded, +) +from sglang.srt.layers.quantization.fp8_utils import ( + _check_cutlass_block_fp8_hardware_support, + cutlass_w8a8_block_fp8_linear_with_fallback, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-large") + +_FP8_MAX = torch.finfo(fp8_dtype).max +_BLOCK = 128 +# Cover M == 1 (greedy decode), small unaligned M (speculative draft tokens), +# the 4-row alignment boundary, and a large aligned batch. +_M_VALUES = [1, 2, 3, 4, 5, 7, 13, 16, 31, 64, 256] + + +def _quant_weight_blockwise(weight_bf16: torch.Tensor, block: int = _BLOCK): + """Block-quantize a (N, K) bf16 weight to fp8 with (N//block, K//block) fp32 scales.""" + n, k = weight_bf16.shape + assert n % block == 0 and k % block == 0 + w = weight_bf16.float().reshape(n // block, block, k // block, block) + amax = w.abs().amax(dim=(1, 3)).clamp(min=1e-12) # (N//block, K//block) + scale = amax / _FP8_MAX + wq = (w / scale[:, None, :, None]).clamp(-_FP8_MAX, _FP8_MAX).to(fp8_dtype) + return wq.reshape(n, k), scale.to(torch.float32) + + +def _legacy_cutlass_linear(x_2d, weight, weight_scale): + """The pre-optimization path: unpadded quant, wrapper pads internally.""" + from sgl_kernel import fp8_blockwise_scaled_mm + + q_input, x_scale = per_token_group_quant_fp8(x_2d, _BLOCK, column_major_scales=True) + return fp8_blockwise_scaled_mm( + q_input, weight.T, x_scale, weight_scale.T, out_dtype=x_2d.dtype + ) + + +@unittest.skipUnless( + _check_cutlass_block_fp8_hardware_support(), + "cutlass block FP8 requires Hopper (SM90) or newer", +) +class TestFP8BlockwiseRowPadding(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.K = 512 + cls.N = 256 + torch.manual_seed(0) + + def test_quant_buffers_row_aligned(self): + """Row-padded quant returns 4-aligned, M-major buffers whose live rows + match the legacy column-major quant bit-for-bit.""" + for m in _M_VALUES: + x = torch.randn(m, self.K, device="cuda", dtype=torch.bfloat16) * 0.1 + xq, xs = sglang_per_token_group_quant_fp8_row_padded(x, _BLOCK) + m_pad = (m + 3) // 4 * 4 + + self.assertEqual(xq.shape, (m_pad, self.K), f"M={m}") + self.assertEqual(xs.shape[0], m_pad, f"M={m}") + # scales_a must stay M-major (stride(0) == 1) for the kernel contract. + self.assertEqual(xs.stride(0), 1, f"M={m}") + + xq_ref, xs_ref = per_token_group_quant_fp8( + x, _BLOCK, column_major_scales=True + ) + self.assertEqual(xq_ref.shape, (m, self.K), f"M={m}") + # Live rows are produced by the same kernel, so they must be identical. + self.assertTrue( + torch.equal(xq[:m].view(torch.uint8), xq_ref.view(torch.uint8)), + f"quantized activation mismatch at M={m}", + ) + torch.testing.assert_close(xs[:m], xs_ref, atol=0.0, rtol=0.0) + + def test_gemm_bit_exact_vs_legacy(self): + """The full linear (row-padded) is bit-identical to the legacy unpadded GEMM.""" + weight_bf16 = ( + torch.randn(self.N, self.K, device="cuda", dtype=torch.bfloat16) * 0.1 + ) + weight, weight_scale = _quant_weight_blockwise(weight_bf16) + + for m in _M_VALUES: + x = torch.randn(m, self.K, device="cuda", dtype=torch.bfloat16) * 0.1 + + out_ref = _legacy_cutlass_linear(x, weight, weight_scale) + out_new = cutlass_w8a8_block_fp8_linear_with_fallback( + input=x, + weight=weight, + block_size=[_BLOCK, _BLOCK], + weight_scale=weight_scale, + ) + + self.assertEqual(out_new.shape, (m, self.N), f"M={m}") + self.assertTrue( + torch.equal(out_ref, out_new), + f"row-padded GEMM differs from legacy at M={m}: " + f"max_abs_diff={(out_ref.float() - out_new.float()).abs().max().item()}", + ) + + def test_linear_matches_bf16_reference(self): + """Sanity: the FP8 linear stays close to a bf16 reference matmul.""" + weight_bf16 = ( + torch.randn(self.N, self.K, device="cuda", dtype=torch.bfloat16) * 0.1 + ) + weight, weight_scale = _quant_weight_blockwise(weight_bf16) + + for m in [1, 5, 64]: + x = torch.randn(m, self.K, device="cuda", dtype=torch.bfloat16) * 0.1 + ref = (x.float() @ weight_bf16.float().T).to(torch.bfloat16) + out = cutlass_w8a8_block_fp8_linear_with_fallback( + input=x, + weight=weight, + block_size=[_BLOCK, _BLOCK], + weight_scale=weight_scale, + ) + torch.testing.assert_close(out, ref, atol=0.5, rtol=0.1) + + +if __name__ == "__main__": + unittest.main() From 46618775eb999e7c7f684629c9be8092056251e9 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Tue, 14 Jul 2026 23:29:38 +0000 Subject: [PATCH 09/25] x --- test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py b/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py index 5f46b4cfee12..7fda67bda8e2 100644 --- a/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py +++ b/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py @@ -34,7 +34,7 @@ ) SERVER_LAUNCH_TIMEOUT = 5400 GSM8K_ACCURACY_THRESHOLD = 0.92 -AVG_SPEC_ACCEPT_LENGTH_THRESHOLD = 4.0 +AVG_SPEC_ACCEPT_LENGTH_THRESHOLD = 3.0 class TestDeepseekV4DSparkUnifiedKVGSM8K(CustomTestCase): From 0e6fcb1ad1588b5d5223dc226dfd7c369952ebee Mon Sep 17 00:00:00 2001 From: At1a8 Date: Wed, 15 Jul 2026 00:14:22 +0000 Subject: [PATCH 10/25] Revert "Revert "Delete CUTLASS FP8 blockwise for SM90 and SM100, move SM120 to JIT and add SwapAB (#30438)"" This reverts commit ac2470d92a145deb4ac3e4b901d8d83378a74404. --- .../fp8_blockwise_scaled_mm_entry.cuh | 25 + .../fp8_blockwise_scaled_mm_sm120.cuh | 502 +++++++++++++++++ .../sglang/jit_kernel/fp8_blockwise_gemm.py | 93 ++++ .../jit_kernel/include/sgl_kernel/utils.cuh | 16 + .../srt/layers/quantization/fp8_utils.py | 29 +- scripts/ci/cuda/warmup_deep_gemm.py | 2 +- sgl-kernel/CMakeLists.txt | 1 - .../benchmark/bench_fp8_blockwise_gemm.py | 237 -------- sgl-kernel/csrc/common_extension.cc | 5 - .../gemm/fp8_blockwise_gemm_sm90_dispatch.cuh | 197 ------- .../csrc/gemm/fp8_blockwise_gemm_kernel.cu | 522 ------------------ sgl-kernel/include/sgl_kernel_ops.h | 6 - sgl-kernel/python/sgl_kernel/__init__.py | 2 - sgl-kernel/python/sgl_kernel/gemm.py | 10 - .../jit/benchmark/bench_fp8_blockwise_gemm.py | 103 ++++ .../jit}/test_fp8_blockwise_gemm.py | 38 +- .../quant/test_fp8_blockwise_row_padding.py | 136 ----- 17 files changed, 765 insertions(+), 1159 deletions(-) create mode 100644 python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh create mode 100644 python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh create mode 100644 python/sglang/jit_kernel/fp8_blockwise_gemm.py delete mode 100644 sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py delete mode 100644 sgl-kernel/csrc/cutlass_extensions/gemm/fp8_blockwise_gemm_sm90_dispatch.cuh delete mode 100644 sgl-kernel/csrc/gemm/fp8_blockwise_gemm_kernel.cu create mode 100644 test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py rename {sgl-kernel/tests => test/registered/jit}/test_fp8_blockwise_gemm.py (69%) delete mode 100644 test/registered/quant/test_fp8_blockwise_row_padding.py diff --git a/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh b/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh new file mode 100644 index 000000000000..cdbecf8a2895 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh @@ -0,0 +1,25 @@ +/* Copyright 2026 SGLang Team. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "fp8_blockwise_scaled_mm_sm120.cuh" + +void fp8_blockwise_scaled_mm( + tvm::ffi::TensorView out, + tvm::ffi::TensorView mat_a, + tvm::ffi::TensorView mat_b, + tvm::ffi::TensorView scales_a, + tvm::ffi::TensorView scales_b) { + fp8_blockwise_scaled_mm_sm120(out, mat_a, mat_b, scales_a, scales_b); +} diff --git a/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh b/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh new file mode 100644 index 000000000000..8f802d92efbe --- /dev/null +++ b/python/sglang/jit_kernel/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh @@ -0,0 +1,502 @@ +/* Copyright 2026 SGLang Team. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include + +#include +#include + +#include +#include +#include + +using namespace host; + +// clang-format off +#include "cutlass/cutlass.h" +#include "cutlass/detail/blockwise_scale_layout.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/util/packed_stride.hpp" +// clang-format on + +#define CUTLASS_CHECK(status) \ + { \ + cutlass::Status error = status; \ + RuntimeCheck(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \ + } + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED) + +template < + typename OutType, + typename MmaTileShape, + typename PerSmTileShape, + typename EpilogueTileShape, + typename ScalesPerTile, + int TileSizeM_ = 128, + class ClusterShape = Shape<_1, _1, _1>> +void launch_sm120_fp8_blockwise_scaled_mm( + tvm::ffi::TensorView out, + tvm::ffi::TensorView a, + tvm::ffi::TensorView b, + tvm::ffi::TensorView scales_a, + tvm::ffi::TensorView scales_b, + cudaStream_t stream) { + using ElementBlockScale = float; + + // A matrix configuration + using ElementA = cutlass::float_e4m3_t; // Element type for A matrix operand + using LayoutATag = cutlass::layout::RowMajor; // Layout type for A matrix operand + constexpr int AlignmentA = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of A matrix in units of + // elements (up to 16 bytes) + + // B matrix configuration + using ElementB = cutlass::float_e4m3_t; // Element type for B matrix operand + using LayoutBTag = cutlass::layout::ColumnMajor; // Layout type for B matrix operand + constexpr int AlignmentB = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of B matrix in units of + // elements (up to 16 bytes) + + // C/D matrix configuration + using ElementD = OutType; // Element type for D matrix operand + using ElementC = void; // Element type for C matrix operand + using LayoutCTag = cutlass::layout::RowMajor; // Layout type for C matrix operand + using LayoutDTag = cutlass::layout::RowMajor; // Layout type for D matrix operand + constexpr int AlignmentD = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of C matrix in units of + // elements (up to 16 bytes) + constexpr int AlignmentC = + AlignmentD; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes) + + // Kernel functional config + using ElementAccumulator = float; // Element type for internal accumulation + using ArchTag = cutlass::arch::Sm120; // Tag indicating the minimum SM that supports the intended feature + using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag - changed from OpClassBlockScaledTensorOp + + static constexpr int ScaleMsPerTile = size<0>(ScalesPerTile{}); + static constexpr int ScaleGranularityM = size<0>(MmaTileShape{}) / ScaleMsPerTile; + static constexpr int ScaleGranularityN = size<1>(MmaTileShape{}) / size<1>(ScalesPerTile{}); + static constexpr int ScaleGranularityK = size<2>(MmaTileShape{}) / size<2>(ScalesPerTile{}); + + using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig< + ScaleGranularityM, + ScaleGranularityN, + ScaleGranularityK, + cute::UMMA::Major::MN, + cute::UMMA::Major::K>; + // FP8 Block-wise scaling configuration + using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); // Layout type for SFA matrix operand + using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); // Layout type for SFB matrix operand + + constexpr bool kCanUsePingpong = (64 % ScaleGranularityM == 0); + + int m = a.size(0); + int k = a.size(1); + int n = b.size(1); + + auto a_ptr = static_cast(a.data_ptr()); + auto b_ptr = static_cast(b.data_ptr()); + auto c_ptr = static_cast(out.data_ptr()); + + auto scales_a_ptr = static_cast(scales_a.data_ptr()); + auto scales_b_ptr = static_cast(scales_b.data_ptr()); + + LayoutSFA layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1)); + LayoutSFB layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1)); + + auto run_gemm = [&](auto tag) -> cutlass::Status { + using GemmKernel = decltype(tag); + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + Gemm gemm_op; + + using StrideA = typename GemmKernel::StrideA; + using StrideB = typename GemmKernel::StrideB; + using StrideC = typename GemmKernel::StrideD; + + StrideA stride_a = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1)); + StrideB stride_b = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1)); + StrideC stride_c = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1)); + + typename GemmKernel::MainloopArguments mainloop_args{ + a_ptr, stride_a, b_ptr, stride_b, scales_a_ptr, layout_SFA, scales_b_ptr, layout_SFB}; + + typename GemmKernel::EpilogueArguments epilogue_args{{}, c_ptr, stride_c, c_ptr, stride_c}; + epilogue_args.thread.alpha = 1.0f; + + typename Gemm::Arguments args = { + cutlass::gemm::GemmUniversalMode::kGemm, + {m, n, k, 1}, + mainloop_args, + epilogue_args, + }; + + auto can_implement = gemm_op.can_implement(args); + if (can_implement != cutlass::Status::kSuccess) { + return can_implement; + } + + size_t workspace_size = gemm_op.get_workspace_size(args); + auto workspace_tensor = alloc_workspace_tensor(workspace_size, a.device()); + void* workspace = (workspace_size == 0) ? nullptr : workspace_tensor.data_ptr(); + + auto init_status = gemm_op.initialize(args, workspace, stream); + if (init_status != cutlass::Status::kSuccess) { + return init_status; + } + + return gemm_op.run(stream); + }; + + using CooperativeCollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + PerSmTileShape, + ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, + ElementAccumulator, + ElementC, + LayoutCTag, + AlignmentC, + ElementD, + LayoutDTag, + AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; + + using CooperativeStageCount = cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CooperativeCollectiveEpilogue::SharedStorage))>; + + using CooperativeCollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + ElementA, + cute::tuple, + AlignmentA, + ElementB, + cute::tuple, + AlignmentB, + ElementAccumulator, + MmaTileShape, + ClusterShape, + CooperativeStageCount, + cutlass::gemm::KernelScheduleSm120Blockwise>::CollectiveOp; + + using CooperativeGemmKernelStreamK = cutlass::gemm::kernel::GemmUniversal< + Shape, + CooperativeCollectiveMainloop, + CooperativeCollectiveEpilogue, + cutlass::gemm::StreamKScheduler>; + using CooperativeGemmKernelVoid = cutlass::gemm::kernel:: + GemmUniversal, CooperativeCollectiveMainloop, CooperativeCollectiveEpilogue, void>; + + auto run_cooperative = [&]() -> cutlass::Status { + static const uint32_t kNumSM = host::runtime::get_sm_count(a.device().device_id); + constexpr int kTileM = size<0>(MmaTileShape{}); + constexpr int kTileN = size<1>(MmaTileShape{}); + uint64_t tiles = static_cast((m + kTileM - 1) / kTileM) * ((n + kTileN - 1) / kTileN); + uint32_t last_wave = static_cast(tiles % kNumSM); + if (last_wave == 0) last_wave = kNumSM; + float waste = 1.0f - static_cast(last_wave) / static_cast(kNumSM); + return (waste > 0.5f) ? run_gemm(CooperativeGemmKernelStreamK{}) : run_gemm(CooperativeGemmKernelVoid{}); + }; + + cutlass::Status status = cutlass::Status::kSuccess; + if constexpr (kCanUsePingpong) { + using PingpongMmaTileShape_MNK = Shape<_64, _128, _128>; + using PingpongCollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + PerSmTileShape, + ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, + ElementAccumulator, + ElementC, + LayoutCTag, + AlignmentC, + ElementD, + LayoutDTag, + AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; + + using PingpongStageCount = cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename PingpongCollectiveEpilogue::SharedStorage))>; + + using PingpongCollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + ElementA, + cute::tuple, + AlignmentA, + ElementB, + cute::tuple, + AlignmentB, + ElementAccumulator, + PingpongMmaTileShape_MNK, + ClusterShape, + PingpongStageCount, + cutlass::gemm::KernelTmaWarpSpecializedBlockwisePingpongSm120>::CollectiveOp; + + using PingpongGemmKernel = cutlass::gemm::kernel:: + GemmUniversal, PingpongCollectiveMainloop, PingpongCollectiveEpilogue, void>; + + if (m <= 64) { + status = run_gemm(PingpongGemmKernel{}); + if (status != cutlass::Status::kSuccess) { + status = run_cooperative(); + } + } else { + status = run_cooperative(); + } + } else { + status = run_cooperative(); + } + + CUTLASS_CHECK(status); +} + +// Transposed GEMM D^T = Wgemm(weight, activation): puts tokens on the N axis. +template < + typename OutType, + typename MmaTileShape, + typename PerSmTileShape, + typename EpilogueTileShape, + typename ScalesPerTile, + class ClusterShape = Shape<_1, _1, _1>> +void launch_sm120_fp8_blockwise_scaled_mm_swapab( + tvm::ffi::TensorView out, + tvm::ffi::TensorView a, + tvm::ffi::TensorView b, + tvm::ffi::TensorView scales_a, + tvm::ffi::TensorView scales_b, + cudaStream_t stream) { + using ElementBlockScale = float; + + using ElementA = cutlass::float_e4m3_t; // A' = weight + using LayoutATag = cutlass::layout::RowMajor; // weight [N, K] is row-major + constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; + + using ElementB = cutlass::float_e4m3_t; // B' = activation + using LayoutBTag = cutlass::layout::ColumnMajor; // activation as [K, M] column-major + constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; + + using ElementD = OutType; + using ElementC = void; + using LayoutCTag = cutlass::layout::ColumnMajor; // D' = out^T is column-major + using LayoutDTag = cutlass::layout::ColumnMajor; + constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + constexpr int AlignmentC = AlignmentD; + + using ElementAccumulator = float; + using ArchTag = cutlass::arch::Sm120; + using OperatorClass = cutlass::arch::OpClassTensorOp; + + static constexpr int ScaleMsPerTile = size<0>(ScalesPerTile{}); + static constexpr int ScaleGranularityM = size<0>(MmaTileShape{}) / ScaleMsPerTile; + static constexpr int ScaleGranularityN = size<1>(MmaTileShape{}) / size<1>(ScalesPerTile{}); + static constexpr int ScaleGranularityK = size<2>(MmaTileShape{}) / size<2>(ScalesPerTile{}); + + // Operands are swapped, so the scale majors swap relative to the non-swap path: + // SFA (weight) is K-major; SFB (per-token activation) is MN-major. + using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig< + ScaleGranularityM, + ScaleGranularityN, + ScaleGranularityK, + cute::UMMA::Major::K, + cute::UMMA::Major::MN>; + using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); + using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); + + int m = a.size(0); // original tokens -> swapped N' + int k = a.size(1); + int n = b.size(1); // original weight cols -> swapped M' + + auto weight_ptr = static_cast(b.data_ptr()); + auto act_ptr = static_cast(a.data_ptr()); + auto c_ptr = static_cast(out.data_ptr()); + auto weight_scale_ptr = static_cast(scales_b.data_ptr()); + auto act_scale_ptr = static_cast(scales_a.data_ptr()); + + // Swapped problem shape (M', N', K) = (n, m, k). + LayoutSFA layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(make_shape(n, m, k, 1)); + LayoutSFB layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1)); + + auto run_gemm = [&](auto tag) -> cutlass::Status { + using GemmKernel = decltype(tag); + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + Gemm gemm_op; + + using StrideA = typename GemmKernel::StrideA; + using StrideB = typename GemmKernel::StrideB; + using StrideC = typename GemmKernel::StrideD; + + StrideA stride_a = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(n, k, 1)); + StrideB stride_b = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(m, k, 1)); + StrideC stride_c = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(n, m, 1)); + + typename GemmKernel::MainloopArguments mainloop_args{ + weight_ptr, stride_a, act_ptr, stride_b, weight_scale_ptr, layout_SFA, act_scale_ptr, layout_SFB}; + + typename GemmKernel::EpilogueArguments epilogue_args{{}, c_ptr, stride_c, c_ptr, stride_c}; + epilogue_args.thread.alpha = 1.0f; + + typename Gemm::Arguments args = { + cutlass::gemm::GemmUniversalMode::kGemm, + {n, m, k, 1}, + mainloop_args, + epilogue_args, + }; + + auto can_implement = gemm_op.can_implement(args); + if (can_implement != cutlass::Status::kSuccess) { + return can_implement; + } + + size_t workspace_size = gemm_op.get_workspace_size(args); + auto workspace_tensor = alloc_workspace_tensor(workspace_size, a.device()); + void* workspace = (workspace_size == 0) ? nullptr : workspace_tensor.data_ptr(); + + auto init_status = gemm_op.initialize(args, workspace, stream); + if (init_status != cutlass::Status::kSuccess) { + return init_status; + } + + return gemm_op.run(stream); + }; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + PerSmTileShape, + ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, + ElementAccumulator, + ElementC, + LayoutCTag, + AlignmentC, + ElementD, + LayoutDTag, + AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; + + using StageCount = cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + ElementA, + cute::tuple, + AlignmentA, + ElementB, + cute::tuple, + AlignmentB, + ElementAccumulator, + MmaTileShape, + ClusterShape, + StageCount, + cutlass::gemm::KernelScheduleSm120Blockwise>::CollectiveOp; + + using GemmKernel = + cutlass::gemm::kernel::GemmUniversal, CollectiveMainloop, CollectiveEpilogue, void>; + + CUTLASS_CHECK(run_gemm(GemmKernel{})); +} + +// swapAB (tile N=32) beats the non-swap 128x128 path for M<=64 or M%4!=0 +// (cold-L2 CUPTI benchmarks, up to ~1.2x); tile N=16 is unsupported by the +// SM120 blockwise collective (needs EPI_TILE_N=32 | CTA_N and B LDSM N>=32). +template +void sm120_fp8_blockwise_dispatch_shape( + tvm::ffi::TensorView out, + tvm::ffi::TensorView a, + tvm::ffi::TensorView b, + tvm::ffi::TensorView scales_a, + tvm::ffi::TensorView scales_b, + cudaStream_t stream) { + const int m = a.size(0); + using EpilogueTileShape = Shape<_128, _64>; + if (m <= 64 || (m % 4 != 0)) { + launch_sm120_fp8_blockwise_scaled_mm_swapab< + OutType, + Shape<_128, _32, _128>, + Shape<_128, _32, _128>, + EpilogueTileShape, + Shape<_1, _32, _1>>(out, a, b, scales_a, scales_b, stream); + return; + } + + using MmaTileShape = Shape<_128, _128, _128>; + using PerSmTileShape = Shape<_128, _128, _128>; + using ScalesPerTile = Shape<_128, _1, _1>; + launch_sm120_fp8_blockwise_scaled_mm( + out, a, b, scales_a, scales_b, stream); +} + +inline void fp8_blockwise_scaled_mm_sm120( + tvm::ffi::TensorView out, + tvm::ffi::TensorView mat_a, + tvm::ffi::TensorView mat_b, + tvm::ffi::TensorView scales_a, + tvm::ffi::TensorView scales_b) { + RuntimeCheck(mat_a.device().device_type == kDLCUDA, "mat_a must be a CUDA tensor"); + RuntimeCheck(mat_b.device().device_type == kDLCUDA, "mat_b must be a CUDA tensor"); + + RuntimeCheck(mat_a.dim() == 2, "mat_a must be a 2D tensor"); + RuntimeCheck(mat_b.dim() == 2, "mat_b must be a 2D tensor"); + RuntimeCheck(mat_a.stride(1) == 1, "mat_a must be a row major tensor"); + RuntimeCheck(mat_b.stride(0) == 1, "mat_b must be a column major tensor"); + RuntimeCheck(mat_a.size(1) == mat_b.size(0), "mat_a and mat_b shapes cannot be multiplied"); + + RuntimeCheck( + (mat_a.size(1) * (mat_a.dtype().bits / 8)) % 16 == 0, "mat_a must be multiple of 16 bytes for memory alignment"); + RuntimeCheck( + (mat_b.size(0) * (mat_b.dtype().bits / 8)) % 16 == 0, "mat_b must be multiple of 16 bytes for memory alignment"); + RuntimeCheck(host::is_type(mat_a.dtype()), "mat_a must be Float8_e4m3fn"); + RuntimeCheck(host::is_type(mat_b.dtype()), "mat_b must be Float8_e4m3fn"); + + RuntimeCheck(mat_a.size(0) == scales_a.size(0), "size of scales_a is not matched"); + RuntimeCheck(mat_a.size(1) / 128 == scales_a.size(1), "size of scales_a is not matched"); + RuntimeCheck(mat_b.size(0) / 128 == scales_b.size(0), "size of scales_b is not matched"); + RuntimeCheck(mat_b.size(1) / 128 == scales_b.size(1), "size of scales_b is not matched"); + RuntimeCheck(host::is_type(scales_a.dtype()), "scales_a must be Float32"); + RuntimeCheck(host::is_type(scales_b.dtype()), "scales_b must be Float32"); + + RuntimeCheck( + (out.size(1) * (out.dtype().bits / 8)) % 16 == 0, "out must be multiple of 16 bytes for memory alignment"); + + const cudaStream_t stream = LaunchKernel::resolve_device(mat_a.device()); + + if (host::is_type(out.dtype())) { + sm120_fp8_blockwise_dispatch_shape(out, mat_a, mat_b, scales_a, scales_b, stream); + } else if (host::is_type(out.dtype())) { + sm120_fp8_blockwise_dispatch_shape(out, mat_a, mat_b, scales_a, scales_b, stream); + } else { + Panic("out_dtype must be Half or BFloat16"); + } +} + +#endif // defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED) diff --git a/python/sglang/jit_kernel/fp8_blockwise_gemm.py b/python/sglang/jit_kernel/fp8_blockwise_gemm.py new file mode 100644 index 000000000000..49b4c9606e25 --- /dev/null +++ b/python/sglang/jit_kernel/fp8_blockwise_gemm.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING + +import torch + +from sglang.jit_kernel.utils import cache_once, load_jit, override_jit_cuda_arch +from sglang.kernel_api_logging import debug_kernel_api +from sglang.srt.utils.common import is_sm120_supported +from sglang.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +def _fp8_blockwise_cuda_flags() -> list[str]: + return [ + "-DNDEBUG", + "-DCUTE_USE_PACKED_TUPLE=1", + "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1", + "-DCUTLASS_VERSIONS_GENERATED", + "-DCUTLASS_TEST_LEVEL=0", + "-DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1", + "-DCUTLASS_DEBUG_TRACE_LEVEL=0", + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + ] + + +@contextmanager +def _fp8_blockwise_arch_env(): + if not is_sm120_supported(): + raise RuntimeError( + "fp8_blockwise_scaled_mm JIT kernel requires SM120 (Blackwell)." + ) + major, minor = torch.cuda.get_device_capability() + # sm_*a target (e.g. sm_120a) required, not plain sm_120. + with override_jit_cuda_arch(major, minor, suffix="a"): + yield + + +@cache_once +def _jit_fp8_blockwise_module() -> Module: + """Compile and cache the SM120 fp8 blockwise GEMM module (handles fp16 + bf16).""" + with _fp8_blockwise_arch_env(): + return load_jit( + "fp8_blockwise_scaled_mm", + cuda_files=["gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh"], + cuda_wrappers=[ + ("fp8_blockwise_scaled_mm", "fp8_blockwise_scaled_mm"), + ], + extra_dependencies=["cutlass"], + extra_cuda_cflags=_fp8_blockwise_cuda_flags(), + ) + + +@register_custom_op( + op_name="fp8_blockwise_scaled_mm", + mutates_args=["out"], +) +def _fp8_blockwise_scaled_mm_custom_op( + out: torch.Tensor, + mat_a: torch.Tensor, + mat_b: torch.Tensor, + scales_a: torch.Tensor, + scales_b: torch.Tensor, +) -> None: + module = _jit_fp8_blockwise_module() + module.fp8_blockwise_scaled_mm(out, mat_a, mat_b, scales_a, scales_b) + + +@debug_kernel_api +def fp8_blockwise_scaled_mm( + mat_a: torch.Tensor, + mat_b: torch.Tensor, + scales_a: torch.Tensor, + scales_b: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + """FP8 e4m3 block-wise scaled matmul on SM120.""" + assert out_dtype in ( + torch.float16, + torch.bfloat16, + ), f"out_dtype must be Half or BFloat16, got {out_dtype}" + + out = torch.empty( + (mat_a.shape[0], mat_b.shape[1]), + dtype=out_dtype, + device=mat_a.device, + ) + _fp8_blockwise_scaled_mm_custom_op(out, mat_a, mat_b, scales_a, scales_b) + return out diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh index bd2d9ab52caa..c5681b933ad9 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh @@ -15,6 +15,7 @@ #pragma once +#include #include #include @@ -238,6 +239,21 @@ inline void RuntimeDeviceCheck(DebugInfo location = {}) { return RuntimeDeviceCheck(::cudaGetLastError(), location); } +inline int getSMVersion(int device_id) { + int sm_major = 0; + int sm_minor = 0; + RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id)); + RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id)); + return sm_major * 10 + sm_minor; +} + +inline auto alloc_workspace_tensor(size_t required_bytes, DLDevice device) -> tvm::ffi::Tensor { + if (required_bytes == 0) return {}; + DLDataType u8 = {kDLUInt, 8, 1}; + int64_t shape[] = {static_cast(required_bytes)}; + return ffi::empty(tvm::ffi::ShapeView(shape, 1), u8, device); +} + /** * \brief Kernel launcher with automatic stream resolution and PDL support. * diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 31cffa226579..accdd4a4a747 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -158,8 +158,9 @@ def use_aiter_triton_gemm_w8a8_tuned_gfx950(n: int, k: int) -> bool: if _is_cuda: - from sgl_kernel import fp8_blockwise_scaled_mm, fp8_scaled_mm + from sgl_kernel import fp8_scaled_mm + from sglang.jit_kernel.fp8_blockwise_gemm import fp8_blockwise_scaled_mm from sglang.srt.utils.patch_torch import register_fake_if_exists @register_fake_if_exists("sgl_kernel::fp8_scaled_mm") @@ -169,13 +170,6 @@ def _fp8_scaled_mm_abstract(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=No N = mat_b.shape[-1] return mat_a.new_empty((M, N), dtype=out_dtype) - @register_fake_if_exists("sgl_kernel::fp8_blockwise_scaled_mm") - def _fp8_blockwise_scaled_mm_abstract(mat_a, mat_b, scales_a, scales_b, out_dtype): - # mat_a: [M, K], mat_b: [K, N] or [N, K] depending on callsite layout; output is [M, N]. - M = mat_a.shape[-2] - N = mat_b.shape[-1] - return mat_a.new_empty((M, N), dtype=out_dtype) - use_triton_w8a8_fp8_kernel = get_bool_env_var("USE_TRITON_W8A8_FP8_KERNEL") @@ -274,11 +268,6 @@ def is_aiter(self) -> bool: FP8_GEMM_RUNNER_BACKEND: Fp8GemmRunnerBackend | None = None -def _check_cutlass_block_fp8_hardware_support() -> bool: - """Return True if CUTLASS block FP8 is supported (Hopper or newer with CUDA 12.0+).""" - return is_sm90_supported() or is_blackwell_supported() - - if is_blackwell_supported() and is_flashinfer_available(): from flashinfer import SfLayout from flashinfer import bmm_fp8 as _raw_flashinfer_bmm_fp8 @@ -541,11 +530,10 @@ def _dispatch_explicit_backend(backend: Fp8GemmRunnerBackend) -> Callable: return flashinfer_deepgemm_w8a8_block_fp8_linear_with_fallback elif backend.is_cutlass(): - if not _check_cutlass_block_fp8_hardware_support(): + if not is_sm120_supported(): raise RuntimeError( - "CUTLASS block FP8 requested via --fp8-gemm-backend=cutlass, " - "but hardware does not support it. CUTLASS block FP8 requires " - "Hopper (SM90+) GPUs with CUDA 12.0+." + "--fp8-gemm-backend=cutlass is deprecated on this hardware. " + "Please switch to DeepGEMM or FlashInfer TRTLLM on SM90/SM100." ) return cutlass_w8a8_block_fp8_linear_with_fallback @@ -579,7 +567,7 @@ def _dispatch_auto_backend() -> Callable: # Priority order for auto selection: # 1. DeepGEMM (if enabled and available) # 2. FlashInfer TRTLLM (if Blackwell GPU and FlashInfer available) - # 3. CUTLASS (if Hopper+ GPU and CUDA 12.0+) + # 3. CUTLASS (if SM120 GPU and CUDA 12.8+) # 4. AITER (if AMD GPU with AITER enabled) # 5. Triton (fallback) @@ -587,7 +575,7 @@ def _dispatch_auto_backend() -> Callable: return deepgemm_w8a8_block_fp8_linear_with_fallback elif is_blackwell_supported() and is_flashinfer_available(): return flashinfer_gemm_w8a8_block_fp8_linear_with_fallback - elif _check_cutlass_block_fp8_hardware_support(): + elif is_sm120_supported(): return cutlass_w8a8_block_fp8_linear_with_fallback elif _use_aiter: return aiter_w8a8_block_fp8_linear @@ -601,8 +589,7 @@ def initialize_fp8_gemm_config(server_args: ServerArgs) -> None: backend = server_args.fp8_gemm_runner_backend if backend == "auto" and is_sm120_supported(): - # TODO(brayden): Verify if CUTLASS can be set by default once SwapAB is supported - backend = "triton" + backend = "cutlass" backend = Fp8GemmRunnerBackend(backend) diff --git a/scripts/ci/cuda/warmup_deep_gemm.py b/scripts/ci/cuda/warmup_deep_gemm.py index 58b7c752c22b..270c2e0bd23b 100644 --- a/scripts/ci/cuda/warmup_deep_gemm.py +++ b/scripts/ci/cuda/warmup_deep_gemm.py @@ -115,7 +115,7 @@ def compute_deepseek_v2v3_shapes(config, tp): Shape derivation based on: - MoE: python/sglang/srt/layers/moe/fused_moe_triton/layer.py - MLA: python/sglang/srt/models/deepseek_v2.py - - FP8: python/sglang/srt/layers/quantization/fp8_kernel.py + - FP8: python/sglang/kernels/ops/quantization/fp8_kernel.py """ shapes = [] diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index 3c3a9f41641e..dc9b4ca27f6e 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -265,7 +265,6 @@ set(SOURCES "csrc/gemm/awq_kernel.cu" "csrc/gemm/bmm_fp8.cu" "csrc/gemm/dsv3_fused_a_gemm.cu" - "csrc/gemm/fp8_blockwise_gemm_kernel.cu" "csrc/gemm/fp8_gemm_kernel.cu" "csrc/gemm/int8_gemm_kernel.cu" "csrc/gemm/per_token_group_quant_8bit.cu" diff --git a/sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py b/sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py deleted file mode 100644 index f05687261890..000000000000 --- a/sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py +++ /dev/null @@ -1,237 +0,0 @@ -import argparse -import copy -import itertools -import os - -import deep_gemm -import torch -import triton -from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor -from sgl_kernel import fp8_blockwise_scaled_mm - -from sglang.utils import is_in_ci - -# Optional vLLM import -try: - from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm - - VLLM_AVAILABLE = True -except ImportError: - vllm_scaled_mm = None - VLLM_AVAILABLE = False - -from sglang.kernels.ops.quantization.fp8_kernel import ( - w8a8_block_fp8_matmul_triton as w8a8_block_fp8_matmul, -) - -IS_CI = is_in_ci() - - -def get_weight_shapes(args): - models_tps = list(itertools.product(args.models, args.tp_sizes)) - # NOTE(HandH1998): The weight shapes only works for DeepSeek-V3. Modify them, if you tune for another different model. - # cannot TP - total = [ - (512 + 64, 7168), - ((128 + 64) * 128, 7168), - (128 * (128 + 128), 512), - (7168, 16384), - (7168, 18432), - ] - # N can TP - n_tp = [ - (18432 * 2, 7168), - ((128 + 64) * 128, 7168), - (128 * (128 + 128), 512), - (24576, 1536), - (4096, 7168), - ] - # K can TP - k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] - # only support Deepseek-V3 - SUPPORT_MODEL = ["deepseek-ai/DeepSeek-V3"] - - weight_shapes = [] - for model, tp_size in models_tps: - assert model in SUPPORT_MODEL - for t in total: - new_t = [t[0], t[1], model] - weight_shapes.append(new_t) - for n_t in n_tp: - new_t = [n_t[0] // tp_size, n_t[1], model] - weight_shapes.append(new_t) - for k_t in k_tp: - new_t = [k_t[0], k_t[1] // tp_size, model] - weight_shapes.append(new_t) - return weight_shapes - - -def cdiv(a: int, b: int) -> int: - """Ceiling division.""" - return -(a // -b) - - -def fp8_gemm_deepgemm( - x_fp8: torch.Tensor, - x_scale: torch.Tensor, - y_fp8: torch.Tensor, - y_scale: torch.Tensor, - m: int, - n: int, - k: int, -): - """DeepGEMM implementation of FP8 GEMM""" - out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) - - # Run DeepGEMM kernel - deep_gemm.fp8_gemm_nt((x_fp8, x_scale), (y_fp8, y_scale), out) - return out - - -def scale_shape(shape, group_shape): - assert len(shape) == len(group_shape) - return tuple(cdiv(shape[i], group_shape[i]) for i in range(len(group_shape))) - - -# CI environment uses simplified parameters -if IS_CI: - batch_sizes = [1, 8] # Simplified for CI -else: - batch_sizes = [1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096] - -# Filter providers based on availability -available_providers = ["sgl-kernel"] -available_names = ["sgl-kernel"] -available_styles = [("orange", "-")] - -if VLLM_AVAILABLE: - available_providers.insert(0, "vllm") - available_names.insert(0, "vllm") - available_styles.insert(0, ("blue", "-")) - -available_providers.append("triton") -available_names.append("sglang triton") -available_styles.append(("red", "-")) - -# Add deepgemm if available -try: - import deep_gemm - - available_providers.append("deepgemm") - available_names.append("deepgemm") - available_styles.append(("yellow", "-")) -except ImportError: - pass - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["batch_size"], - x_vals=batch_sizes, - x_log=False, - line_arg="provider", - line_vals=available_providers, - line_names=available_names, - styles=available_styles, - ylabel="GB/s", - plot_name="fp8 blockwise scaled matmul", - args={}, - ) -) -def benchmark(batch_size, provider, N, K): - M = batch_size - fp8_info = torch.finfo(torch.float8_e4m3fn) - fp8_max, fp8_min = fp8_info.max, fp8_info.min - - a_fp32 = (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max - a_fp8 = a_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) - - b_fp32 = (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max - b_fp8 = b_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) - - scale_a_group_shape = (1, 128) - scale_b_group_shape = (128, 128) - scale_a_shape = scale_shape(a_fp8.shape, scale_a_group_shape) - scale_b_shape = scale_shape(b_fp8.shape, scale_b_group_shape) - - scale_a = torch.randn(scale_a_shape, device="cuda", dtype=torch.float32) - scale_b = torch.randn(scale_b_shape, device="cuda", dtype=torch.float32) - - quantiles = [0.5, 0.2, 0.8] - if provider == "sgl-kernel": - scale_a = scale_a.t().contiguous().t() - b_fp8, scale_b = b_fp8.t(), scale_b.t() - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( - lambda: fp8_blockwise_scaled_mm( - a_fp8, b_fp8, scale_a, scale_b, torch.float16 - ), - quantiles=quantiles, - ) - elif provider == "vllm": - if not VLLM_AVAILABLE: - return (0, 0, 0) - scale_a = scale_a.t().contiguous().t() - b_fp8, scale_b = b_fp8.t(), scale_b.t() - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( - lambda: vllm_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, torch.float16), - quantiles=quantiles, - ) - elif provider == "triton": - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( - lambda: w8a8_block_fp8_matmul( - a_fp8, b_fp8, scale_a, scale_b, [128, 128], torch.float16 - ), - quantiles=quantiles, - ) - if provider == "deepgemm": - scale_a_col_major = get_mn_major_tma_aligned_tensor(scale_a.clone()) - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( - lambda: fp8_gemm_deepgemm( - a_fp8, scale_a_col_major, b_fp8, scale_b, M, N, K - ), - quantiles=quantiles, - ) - return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--models", - nargs="+", - type=str, - default=["deepseek-ai/DeepSeek-V3"], - help="List of models to benchmark", - ) - parser.add_argument( - "--tp-sizes", - nargs="+", - type=int, - default=[1], - help="List of tensor parallel sizes", - ) - args = parser.parse_args() - - # Simplify for CI environment - if IS_CI: - args.models = [args.models[0]] # Use only first model - args.tp_sizes = [args.tp_sizes[0]] # Use only first TP size - - NK_model_names = get_weight_shapes(args) - - # Limit iterations in CI - if IS_CI: - NK_model_names = NK_model_names[:2] # Only test first 2 shapes in CI - - for N, K, model_name in NK_model_names: - if N % 128 != 0 or K % 128 != 0: - print(f"Skip {N=}, {K=} now") - continue - print(f"{model_name} N={N} K={K}: ") - benchmark.run( - print_data=True, - N=N, - K=K, - ) - - print("Benchmark finished!") diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index 45d3dfe27b55..134aaf453b1d 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -123,11 +123,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "bias) -> Tensor"); m.impl("fp8_scaled_mm", torch::kCUDA, &fp8_scaled_mm); - m.def( - "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> " - "Tensor"); - m.impl("fp8_blockwise_scaled_mm", torch::kCUDA, &fp8_blockwise_scaled_mm); - m.def( "sgl_per_token_group_quant_8bit(Tensor input, Tensor! output_q, Tensor! output_s, int group_size," " float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()"); diff --git a/sgl-kernel/csrc/cutlass_extensions/gemm/fp8_blockwise_gemm_sm90_dispatch.cuh b/sgl-kernel/csrc/cutlass_extensions/gemm/fp8_blockwise_gemm_sm90_dispatch.cuh deleted file mode 100644 index 05b70c4f26f2..000000000000 --- a/sgl-kernel/csrc/cutlass_extensions/gemm/fp8_blockwise_gemm_sm90_dispatch.cuh +++ /dev/null @@ -1,197 +0,0 @@ -// Adapted from -// https://github.com/vllm-project/vllm/blob/main/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh -#pragma once - -#include "cute/tensor.hpp" -#include "cutlass/cutlass.h" -#include "cutlass/epilogue/collective/collective_builder.hpp" -#include "cutlass/epilogue/dispatch_policy.hpp" -#include "cutlass/gemm/collective/collective_builder.hpp" -#include "cutlass/gemm/device/gemm_universal_adapter.h" -#include "cutlass/gemm/dispatch_policy.hpp" -#include "cutlass/gemm/kernel/gemm_universal.hpp" -#include "cutlass/gemm/kernel/tile_scheduler_params.h" -#include "cutlass/numeric_types.h" -#include "cutlass/tensor_ref.h" -#include "cutlass_extensions/common.hpp" -#include "cutlass_extensions/gemm/cutlass_gemm_caller.cuh" -#include "cutlass_extensions/gemm/dispatch_policy.hpp" - -using namespace cute; - -template < - typename SchedulerType, - typename OutType, - int GroupSizeM_, - int GroupSizeN_, - int GroupSizeK_, - int TileSizeM_ = 128, - class ClusterShape = Shape<_1, _2, _1>> -struct cutlass_3x_gemm_fp8_blockwise { - using GroupSizeM = Int; - using GroupSizeN = Int; - using GroupSizeK = Int; - using TileSizeM = Int; - - static_assert(TileSizeM_ % GroupSizeM_ == 0, "TileSizeM must be a multiple of GroupSizeM"); - - using ElementAB = cutlass::float_e4m3_t; - - // A matrix configuration - using ElementA = ElementAB; - using LayoutA = cutlass::layout::RowMajor; - static constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; - - // B matrix configuration - using ElementB = ElementAB; - using LayoutB = cutlass::layout::ColumnMajor; - static constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; - - // C/D matrix configuration - using ElementC = void; - using LayoutC = cutlass::layout::RowMajor; - static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; - - using ElementD = OutType; - using LayoutD = cutlass::layout::RowMajor; - static constexpr int AlignmentD = AlignmentC; - - using ScaleTileShape = Shape<_1, _128, _128>; - using ScaleConfig = decltype(cutlass::detail::sm90_trivial_blockwise_scale_config(ScaleTileShape{})); - using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); - using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); - - // Multiply-accumulate blocking/pipelining details - using ElementAccumulator = float; // Element type for internal accumulation - using ElementCompute = float; // Element type for compute - using TileShape = Shape; // Threadblock-level tile size - - using ArchTag = cutlass::arch::Sm90; - using OperatorClass = cutlass::arch::OpClassTensorOp; - using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative; - using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; - using StoreEpilogueCompute = typename cutlass::epilogue::fusion::Sm90EVT; - - using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8Blockwise; - using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - TileShape, - ClusterShape, - EpilogueTileType, - ElementAccumulator, - ElementCompute, - ElementC, - LayoutC, - AlignmentC, - ElementD, - LayoutD, - AlignmentD, - EpilogueSchedule, - StoreEpilogueCompute>::CollectiveOp; - - using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - ElementA, - cute::tuple, - AlignmentA, - ElementB, - cute::tuple, - AlignmentB, - ElementAccumulator, - TileShape, - ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename CollectiveEpilogue::SharedStorage))>, - KernelSchedule>::CollectiveOp; - - using GemmKernel = cutlass::gemm::kernel::GemmUniversal< - Shape, // Indicates ProblemShape - CollectiveMainloop, - CollectiveEpilogue, - SchedulerType>; -}; - -template -void cutlass_gemm_caller_blockwise( - torch::Tensor& out, - torch::Tensor const& a, - torch::Tensor const& b, - torch::Tensor const& a_scales, - torch::Tensor const& b_scales) { - using GemmKernel = typename Gemm::GemmKernel; - using ElementAB = typename Gemm::ElementAB; - using ElementA = ElementAB; - using ElementB = ElementAB; - using ElementD = typename Gemm::ElementD; - using ElementBlockScale = float; - - using ScaleTileShape = Shape<_1, _128, _128>; - using ScaleConfig = decltype(cutlass::detail::sm90_trivial_blockwise_scale_config(ScaleTileShape{})); - using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); - using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); - - int m = a.size(0); - int k = a.size(1); - int n = b.size(1); - - auto a_ptr = static_cast(a.data_ptr()); - auto b_ptr = static_cast(b.data_ptr()); - - auto a_s_ptr = static_cast(a_scales.data_ptr()); - auto b_s_ptr = static_cast(b_scales.data_ptr()); - - using StrideA = typename GemmKernel::StrideA; - using StrideB = typename GemmKernel::StrideB; - using StrideD = typename GemmKernel::StrideD; - using StrideC = typename GemmKernel::StrideC; - - StrideA a_stride = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1)); - StrideB b_stride = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1)); - StrideC c_stride = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1)); - LayoutSFA layout_sfa = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1)); - LayoutSFB layout_sfb = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1)); - - typename GemmKernel::MainloopArguments mainloop_args{ - a_ptr, a_stride, b_ptr, b_stride, a_s_ptr, layout_sfa, b_s_ptr, layout_sfb}; - auto c_ptr = static_cast(out.data_ptr()); - typename GemmKernel::EpilogueArguments epilogue_args{{}, c_ptr, c_stride, c_ptr, c_stride}; - - typename GemmKernel::TileSchedulerArguments scheduler; - - static constexpr bool UsesStreamKScheduler = - cute::is_same_v; - - if constexpr (UsesStreamKScheduler) { - using DecompositionMode = - typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90StreamKParams::DecompositionMode; - using ReductionMode = - typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90StreamKParams::ReductionMode; - - scheduler.decomposition_mode = DecompositionMode::StreamK; - scheduler.reduction_mode = ReductionMode::Nondeterministic; - } - - cutlass_gemm_caller(a.device(), {m, n, k, 1}, mainloop_args, epilogue_args, scheduler); -} - -template -void cutlass_gemm_blockwise_sm90_fp8_dispatch( - torch::Tensor& out, - torch::Tensor const& a, - torch::Tensor const& b, - torch::Tensor const& a_scales, - torch::Tensor const& b_scales) { - auto k = a.size(1); - auto n = b.size(1); - - if (k > 3 * n) { - cutlass_gemm_caller_blockwise>( - out, a, b, a_scales, b_scales); - } else { - cutlass_gemm_caller_blockwise< - cutlass_3x_gemm_fp8_blockwise>( - out, a, b, a_scales, b_scales); - } -} diff --git a/sgl-kernel/csrc/gemm/fp8_blockwise_gemm_kernel.cu b/sgl-kernel/csrc/gemm/fp8_blockwise_gemm_kernel.cu deleted file mode 100644 index cc094de51a60..000000000000 --- a/sgl-kernel/csrc/gemm/fp8_blockwise_gemm_kernel.cu +++ /dev/null @@ -1,522 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "cutlass_extensions/gemm/cutlass_gemm_caller.cuh" -#include "cutlass_extensions/gemm/fp8_blockwise_gemm_sm90_dispatch.cuh" -#include "utils.h" - -using namespace cute; - -template < - typename OutType, - typename MmaTileShape, - typename PerSmTileShape, - typename EpilogueTileShape, - typename ScalesPerTile, - int TileSizeM_ = 128, - class ClusterShape = Shape<_1, _1, _1>> -void launch_sm100_fp8_blockwise_scaled_mm( - torch::Tensor& out, - const torch::Tensor& a, - const torch::Tensor& b, - const torch::Tensor& scales_a, - const torch::Tensor& scales_b) { - static constexpr int ScaleMsPerTile = size<0>(ScalesPerTile{}); - static constexpr int ScaleGranularityM = size<0>(MmaTileShape{}) / ScaleMsPerTile; - static constexpr int ScaleGranularityN = size<1>(MmaTileShape{}) / size<1>(ScalesPerTile{}); - static constexpr int ScaleGranularityK = size<2>(MmaTileShape{}) / size<2>(ScalesPerTile{}); - - using ElementAB = cutlass::float_e4m3_t; - using ElementA = ElementAB; - using ElementB = ElementAB; - using ElementC = void; - using ElementD = OutType; - using LayoutA = cutlass::layout::RowMajor; - using LayoutB = cutlass::layout::ColumnMajor; - using LayoutD = cutlass::layout::RowMajor; - using LayoutC = LayoutD; - // This means both SFA and SFB are column-major. - using ScaleConfig = cutlass::detail::Sm100BlockwiseScaleConfig< - ScaleGranularityM, - ScaleGranularityN, - ScaleGranularityK, - cute::UMMA::Major::MN, - cute::UMMA::Major::K>; - using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); - using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); - - static constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; - static constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; - static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; - static constexpr int AlignmentC = AlignmentD; - - using ElementAccumulator = float; - using ElementBlockScale = float; - using ElementCompute = float; - using ArchTag = cutlass::arch::Sm100; - using OperatorClass = cutlass::arch::OpClassTensorOp; - - using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - ArchTag, - cutlass::arch::OpClassTensorOp, - PerSmTileShape, - ClusterShape, - EpilogueTileShape, - ElementAccumulator, - ElementCompute, - ElementC, - LayoutC, - AlignmentC, - ElementD, - LayoutD, - AlignmentD, - cutlass::epilogue::TmaWarpSpecialized1Sm>::CollectiveOp; - - using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - ElementA, - cute::tuple, - AlignmentA, - ElementB, - cute::tuple, - AlignmentB, - ElementAccumulator, - MmaTileShape, - ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename CollectiveEpilogue::SharedStorage))>, - cutlass::gemm::KernelTmaWarpSpecializedBlockwise1SmSm100>::CollectiveOp; - - using GemmKernel = cutlass::gemm::kernel::GemmUniversal< - Shape, - CollectiveMainloop, - CollectiveEpilogue, - cutlass::gemm::PersistentScheduler>; - using Gemm = cutlass::gemm::device::GemmUniversalAdapter; - - Gemm gemm_op; - - int m = a.size(0); - int k = a.size(1); - int n = b.size(1); - - auto a_ptr = static_cast(a.data_ptr()); - auto b_ptr = static_cast(b.data_ptr()); - auto scales_a_ptr = static_cast(scales_a.data_ptr()); - auto scales_b_ptr = static_cast(scales_b.data_ptr()); - auto c_ptr = static_cast(out.data_ptr()); - - using StrideA = typename GemmKernel::StrideA; - using StrideB = typename GemmKernel::StrideB; - using StrideD = typename GemmKernel::StrideD; - using StrideC = typename GemmKernel::StrideD; - - StrideA a_stride = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1)); - StrideB b_stride = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1)); - StrideC c_stride = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1)); - LayoutSFA layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1)); - LayoutSFB layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1)); - - typename GemmKernel::MainloopArguments mainloop_args{ - a_ptr, a_stride, b_ptr, b_stride, scales_a_ptr, layout_SFA, scales_b_ptr, layout_SFB}; - - typename GemmKernel::EpilogueArguments epilogue_args{{}, c_ptr, c_stride, c_ptr, c_stride}; - epilogue_args.thread.alpha = 1.0f; - - typename GemmKernel::Arguments args = { - cutlass::gemm::GemmUniversalMode::kGemm, {m, n, k, 1}, mainloop_args, epilogue_args}; - - auto can_implement = gemm_op.can_implement(args); - TORCH_CHECK(can_implement == cutlass::Status::kSuccess, cutlassGetStatusString(can_implement)) - - size_t workspace_size = gemm_op.get_workspace_size(args); - cutlass::device_memory::allocation workspace(workspace_size); - - auto init_status = gemm_op.initialize(args, workspace.get()); - TORCH_CHECK(init_status == cutlass::Status::kSuccess, cutlassGetStatusString(init_status)); - - auto stream = at::cuda::getCurrentCUDAStream(a.get_device()); - - auto status = gemm_op.run(stream); - TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)) -} - -template -void sm100_fp8_blockwise_dispatch_shape( - torch::Tensor& out, - const torch::Tensor& a, - const torch::Tensor& b, - const torch::Tensor& scales_a, - const torch::Tensor& scales_b) { - if (a.size(0) <= 128) { - using MmaTileShape = Shape<_64, _128, _128>; - using PerSmTileShape = Shape<_64, _128, _128>; - using EpilogueTileShape = Shape<_64, _64>; - using ScalesPerTile = Shape<_64, _1, _1>; - launch_sm100_fp8_blockwise_scaled_mm( - out, a, b, scales_a, scales_b); - } else { - using MmaTileShape = Shape<_128, _128, _128>; - using PerSmTileShape = Shape<_128, _128, _128>; - using EpilogueTileShape = Shape<_128, _64>; - using ScalesPerTile = Shape<_128, _1, _1>; - launch_sm100_fp8_blockwise_scaled_mm( - out, a, b, scales_a, scales_b); - } -} - -template < - typename OutType, - typename MmaTileShape, - typename PerSmTileShape, - typename EpilogueTileShape, - typename ScalesPerTile, - int TileSizeM_ = 128, - class ClusterShape = Shape<_1, _1, _1>> -void launch_sm120_fp8_blockwise_scaled_mm( - torch::Tensor& out, - const torch::Tensor& a, - const torch::Tensor& b, - const torch::Tensor& scales_a, - const torch::Tensor& scales_b) { - using ElementBlockScale = float; - - // A matrix configuration - using ElementA = cutlass::float_e4m3_t; // Element type for A matrix operand - using LayoutATag = cutlass::layout::RowMajor; // Layout type for A matrix operand - constexpr int AlignmentA = - 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of A matrix in units of - // elements (up to 16 bytes) - - // B matrix configuration - using ElementB = cutlass::float_e4m3_t; // Element type for B matrix operand - using LayoutBTag = cutlass::layout::ColumnMajor; // Layout type for B matrix operand - constexpr int AlignmentB = - 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of B matrix in units of - // elements (up to 16 bytes) - - // C/D matrix configuration - using ElementD = OutType; // Element type for D matrix operand - using ElementC = void; // Element type for C matrix operand - using LayoutCTag = cutlass::layout::RowMajor; // Layout type for C matrix operand - using LayoutDTag = cutlass::layout::RowMajor; // Layout type for D matrix operand - constexpr int AlignmentD = - 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of C matrix in units of - // elements (up to 16 bytes) - constexpr int AlignmentC = - AlignmentD; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes) - - // Kernel functional config - using ElementAccumulator = float; // Element type for internal accumulation - using ArchTag = cutlass::arch::Sm120; // Tag indicating the minimum SM that supports the intended feature - using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag - changed from OpClassBlockScaledTensorOp - - static constexpr int ScaleMsPerTile = size<0>(ScalesPerTile{}); - static constexpr int ScaleGranularityM = size<0>(MmaTileShape{}) / ScaleMsPerTile; - static constexpr int ScaleGranularityN = size<1>(MmaTileShape{}) / size<1>(ScalesPerTile{}); - static constexpr int ScaleGranularityK = size<2>(MmaTileShape{}) / size<2>(ScalesPerTile{}); - - using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig< - ScaleGranularityM, - ScaleGranularityN, - ScaleGranularityK, - cute::UMMA::Major::MN, - cute::UMMA::Major::K>; - // FP8 Block-wise scaling configuration - using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); // Layout type for SFA matrix operand - using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); // Layout type for SFB matrix operand - - constexpr bool kCanUsePingpong = (64 % ScaleGranularityM == 0); - - int m = a.size(0); - int k = a.size(1); - int n = b.size(1); - - auto a_ptr = static_cast(a.data_ptr()); - auto b_ptr = static_cast(b.data_ptr()); - auto c_ptr = static_cast(out.data_ptr()); - - auto scales_a_ptr = static_cast(scales_a.data_ptr()); - auto scales_b_ptr = static_cast(scales_b.data_ptr()); - - LayoutSFA layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1)); - LayoutSFB layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1)); - - auto run_gemm = [&](auto tag) -> cutlass::Status { - using GemmKernel = decltype(tag); - using Gemm = cutlass::gemm::device::GemmUniversalAdapter; - Gemm gemm_op; - - using StrideA = typename GemmKernel::StrideA; - using StrideB = typename GemmKernel::StrideB; - using StrideC = typename GemmKernel::StrideD; - - StrideA stride_a = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1)); - StrideB stride_b = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1)); - StrideC stride_c = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1)); - - typename GemmKernel::MainloopArguments mainloop_args{ - a_ptr, stride_a, b_ptr, stride_b, scales_a_ptr, layout_SFA, scales_b_ptr, layout_SFB}; - - typename GemmKernel::EpilogueArguments epilogue_args{{}, c_ptr, stride_c, c_ptr, stride_c}; - epilogue_args.thread.alpha = 1.0f; - - typename Gemm::Arguments args = { - cutlass::gemm::GemmUniversalMode::kGemm, - {m, n, k, 1}, - mainloop_args, - epilogue_args, - }; - - auto can_implement = gemm_op.can_implement(args); - if (can_implement != cutlass::Status::kSuccess) { - return can_implement; - } - - size_t workspace_size = gemm_op.get_workspace_size(args); - cutlass::device_memory::allocation workspace(workspace_size); - - auto init_status = gemm_op.initialize(args, workspace.get()); - if (init_status != cutlass::Status::kSuccess) { - return init_status; - } - - auto stream = at::cuda::getCurrentCUDAStream(a.get_device()); - return gemm_op.run(stream); - }; - - using CooperativeCollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - PerSmTileShape, - ClusterShape, - cutlass::epilogue::collective::EpilogueTileAuto, - ElementAccumulator, - ElementAccumulator, - ElementC, - LayoutCTag, - AlignmentC, - ElementD, - LayoutDTag, - AlignmentD, - cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; - - using CooperativeStageCount = cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename CooperativeCollectiveEpilogue::SharedStorage))>; - - using CooperativeCollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - ElementA, - cute::tuple, - AlignmentA, - ElementB, - cute::tuple, - AlignmentB, - ElementAccumulator, - MmaTileShape, - ClusterShape, - CooperativeStageCount, - cutlass::gemm::KernelScheduleSm120Blockwise>::CollectiveOp; - - using CooperativeGemmKernel = cutlass::gemm::kernel:: - GemmUniversal, CooperativeCollectiveMainloop, CooperativeCollectiveEpilogue, void>; - - cutlass::Status status = cutlass::Status::kSuccess; - if constexpr (kCanUsePingpong) { - using PingpongMmaTileShape_MNK = Shape<_64, _128, _128>; - using PingpongCollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - PerSmTileShape, - ClusterShape, - cutlass::epilogue::collective::EpilogueTileAuto, - ElementAccumulator, - ElementAccumulator, - ElementC, - LayoutCTag, - AlignmentC, - ElementD, - LayoutDTag, - AlignmentD, - cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; - - using PingpongStageCount = cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename PingpongCollectiveEpilogue::SharedStorage))>; - - using PingpongCollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - ElementA, - cute::tuple, - AlignmentA, - ElementB, - cute::tuple, - AlignmentB, - ElementAccumulator, - PingpongMmaTileShape_MNK, - ClusterShape, - PingpongStageCount, - cutlass::gemm::KernelTmaWarpSpecializedBlockwisePingpongSm120>::CollectiveOp; - - using PingpongGemmKernel = cutlass::gemm::kernel:: - GemmUniversal, PingpongCollectiveMainloop, PingpongCollectiveEpilogue, void>; - - if (m <= 64) { - status = run_gemm(PingpongGemmKernel{}); - if (status != cutlass::Status::kSuccess) { - status = run_gemm(CooperativeGemmKernel{}); - } - } else { - status = run_gemm(CooperativeGemmKernel{}); - } - } else { - status = run_gemm(CooperativeGemmKernel{}); - } - - TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); -} - -template -void sm120_fp8_blockwise_dispatch_shape( - torch::Tensor& out, - const torch::Tensor& a, - const torch::Tensor& b, - const torch::Tensor& scales_a, - const torch::Tensor& scales_b) { - using MmaTileShape = Shape<_128, _128, _128>; - using PerSmTileShape = Shape<_128, _128, _128>; - using EpilogueTileShape = Shape<_128, _64>; - using ScalesPerTile = Shape<_128, _1, _1>; - launch_sm120_fp8_blockwise_scaled_mm( - out, a, b, scales_a, scales_b); -} - -torch::Tensor fp8_blockwise_scaled_mm( - const torch::Tensor& mat_a, - const torch::Tensor& mat_b, - const torch::Tensor& scales_a, - const torch::Tensor& scales_b, - const torch::Dtype& out_dtype) { - TORCH_CHECK(mat_a.is_cuda(), "mat_a must be a CUDA tensor"); - TORCH_CHECK(mat_b.is_cuda(), "mat_b must be a CUDA tensor"); - TORCH_CHECK(mat_a.dim() == 2, "mat_a must be a 2D tensor"); - TORCH_CHECK(mat_b.dim() == 2, "mat_b must be a 2D tensor"); - TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor"); - TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor"); - TORCH_CHECK(mat_a.size(1) == mat_b.size(0), "mat_a and mat_b shapes cannot be multiplied"); - - TORCH_CHECK( - (mat_a.size(1) * mat_a.element_size()) % 16 == 0, "mat_a must be multiple of 16 bytes for memory alignment"); - TORCH_CHECK( - (mat_b.size(0) * mat_b.element_size()) % 16 == 0, "mat_b must be multiple of 16 bytes for memory alignment"); - TORCH_CHECK(mat_a.scalar_type() == torch::kFloat8_e4m3fn, "mat_a must be Float8_e4m3fn"); - TORCH_CHECK(mat_b.scalar_type() == torch::kFloat8_e4m3fn, "mat_b must be Float8_e4m3fn"); - TORCH_CHECK(out_dtype == torch::kHalf || out_dtype == torch::kBFloat16, "out_dtype must be Half or BFloat16"); - - auto is_contiguous_vector = [](const torch::Tensor& t) { - auto t_sizes = t.sizes(); - return t.is_contiguous() && - (t.dim() == 1 || (t.dim() == 2 && *std::min_element(t_sizes.begin(), t_sizes.end()) == 1)); - }; - - TORCH_CHECK(mat_a.size(0) == scales_a.size(0), "size of scales_a is not matched"); - TORCH_CHECK(mat_a.size(1) / 128 == scales_a.size(1), "size of scales_a is not matched"); - TORCH_CHECK(scales_a.stride(0) == 1 || is_contiguous_vector(scales_a), "scales_a must be M major"); - TORCH_CHECK(mat_b.size(0) / 128 == scales_b.size(0), "size of scales_b is not matched"); - TORCH_CHECK(mat_b.size(1) / 128 == scales_b.size(1), "size of scales_b is not matched"); - TORCH_CHECK(scales_b.stride(0) == 1 || is_contiguous_vector(scales_b), "scales_b must be K major"); - TORCH_CHECK(scales_a.scalar_type() == torch::kFloat32, "scales_a must be Float32"); - TORCH_CHECK(scales_b.scalar_type() == torch::kFloat32, "scales_b must be Float32"); - - torch::Tensor out = torch::empty({mat_a.size(0), mat_b.size(1)}, mat_a.options().dtype(out_dtype)); - TORCH_CHECK((out.size(1) * out.element_size()) % 16 == 0, "out must be multiple of 16 bytes for memory alignment"); - - auto sm_version = getSMVersion(); - - int64_t original_rows = mat_a.size(0); - torch::Tensor mat_a_padded = pad_tensor(mat_a, /*alignment=*/4); - torch::Tensor scales_a_padded = pad_tensor(scales_a, /*alignment=*/4, /*col_major=*/true); - torch::Tensor out_padded = torch::empty({mat_a_padded.size(0), mat_b.size(1)}, out.options()); - -#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) -#if defined CUDA_VERSION && CUDA_VERSION >= 12000 - if (sm_version == 90) { - torch::Tensor scales_b_contiguous = scales_b.contiguous(); - if (out_dtype == torch::kBFloat16) { - cutlass_gemm_blockwise_sm90_fp8_dispatch( - out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b_contiguous); - } else { - cutlass_gemm_blockwise_sm90_fp8_dispatch( - out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b_contiguous); - } - return out_padded.slice(0, 0, original_rows); - } -#endif -#endif - -#if defined(CUTLASS_ARCH_MMA_SM100A_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) -#if defined CUDA_VERSION && CUDA_VERSION >= 12080 - if (sm_version == 100 -#if CUDA_VERSION >= 12090 - || sm_version == 103 -#endif - ) { - if (out_dtype == torch::kBFloat16) { - sm100_fp8_blockwise_dispatch_shape( - out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b); - } else { - sm100_fp8_blockwise_dispatch_shape(out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b); - } - return out_padded.slice(0, 0, original_rows); - } -#endif -#endif - -#if defined(CUTLASS_ARCH_MMA_SM120A_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) -#if defined(CUDA_VERSION) && CUDA_VERSION >= 12080 - if (sm_version >= 120) { - if (out_dtype == torch::kBFloat16) { - sm120_fp8_blockwise_dispatch_shape( - out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b); - } else { - sm120_fp8_blockwise_dispatch_shape(out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b); - } - return out_padded.slice(0, 0, original_rows); - } -#endif -#endif - - TORCH_CHECK_NOT_IMPLEMENTED( - false, "No implemented fp8_blockwise_scaled_mm for current compute capability: ", sm_version); -} diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 9a92f38bf6bb..3f9cda5029ee 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -235,12 +235,6 @@ torch::Tensor fp8_scaled_mm( const torch::Tensor& scales_b, const torch::Dtype& out_dtype, const c10::optional& bias); -torch::Tensor fp8_blockwise_scaled_mm( - const torch::Tensor& mat_a, - const torch::Tensor& mat_b, - const torch::Tensor& scales_a, - const torch::Tensor& scales_b, - const torch::Dtype& out_dtype); void sgl_per_token_group_quant_8bit( at::Tensor input, at::Tensor output_q, diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index 199c1f92264d..21f5eb90f487 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -57,7 +57,6 @@ awq_dequantize, bmm_fp8, dsv3_fused_a_gemm, - fp8_blockwise_scaled_mm, fp8_scaled_mm, gptq_gemm, gptq_shuffle, @@ -178,7 +177,6 @@ "fast_topk_transform_ragged_fused", "fast_topk_v2", "fp8_blockwise_scaled_grouped_mm", - "fp8_blockwise_scaled_mm", "fp8_scaled_mm", "fused_add_rmsnorm", "fused_qk_norm_rope", diff --git a/sgl-kernel/python/sgl_kernel/gemm.py b/sgl-kernel/python/sgl_kernel/gemm.py index 1d68cdf94943..56d88ae14d99 100644 --- a/sgl-kernel/python/sgl_kernel/gemm.py +++ b/sgl-kernel/python/sgl_kernel/gemm.py @@ -21,16 +21,6 @@ def int8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None): ) -def fp8_blockwise_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype): - return torch.ops.sgl_kernel.fp8_blockwise_scaled_mm.default( - mat_a, - mat_b, - scales_a, - scales_b, - out_dtype, - ) - - def fp8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None): return torch.ops.sgl_kernel.fp8_scaled_mm.default( mat_a, diff --git a/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py b/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py new file mode 100644 index 000000000000..dcab974f154a --- /dev/null +++ b/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import sys + +import torch +import triton + +from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark +from sglang.jit_kernel.fp8_blockwise_gemm import fp8_blockwise_scaled_mm +from sglang.srt.utils import is_sm120_supported +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=5, + stage="base-b-kernel-benchmark", + runner_config="1-gpu-large", +) + + +def _make_inputs(m: int, n: int, k: int, device: str = "cuda"): + fp8_info = torch.finfo(torch.float8_e4m3fn) + fp8_max, fp8_min = fp8_info.max, fp8_info.min + a_fp32 = (torch.rand(m, k, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max + a_fp8 = a_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) + b_fp32 = (torch.rand(n, k, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max + b_fp8 = b_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn).t() + + scale_a = torch.randn((m, k // 128), device=device, dtype=torch.float32) * 0.001 + scale_b = ( + torch.randn((k // 128, n // 128), device=device, dtype=torch.float32) * 0.001 + ) + scale_a = scale_a.t().contiguous().t() + scale_b = scale_b.t().contiguous().t() + return a_fp8, b_fp8, scale_a, scale_b + + +def _torch_ref(a_fp8, b_fp8, scale_a, scale_b): + def group_broadcast(t, shape): + for i, s in enumerate(shape): + if t.shape[i] != s and t.shape[i] != 1: + assert s % t.shape[i] == 0 + t = ( + t.unsqueeze(i + 1) + .expand(*t.shape[: i + 1], s // t.shape[i], *t.shape[i + 1 :]) + .flatten(i, i + 1) + ) + return t + + sa = group_broadcast(scale_a, a_fp8.shape) + sb = group_broadcast(scale_b, b_fp8.shape) + return torch.mm(sa * a_fp8.to(torch.float32), sb * b_fp8.to(torch.float32)).to( + torch.bfloat16 + ) + + +shape_range = get_benchmark_range( + full_range=[ + (16, 4096, 4096), # swapAB tile N=32 + (64, 4096, 4096), # swapAB tile N=64 + (128, 4096, 4096), # non-swap 128 + (512, 4096, 4096), + (1024, 8192, 4096), + ], + ci_range=[(16, 4096, 4096), (128, 4096, 4096)], +) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["m", "n", "k"], + x_vals=shape_range, + x_log=False, + line_arg="provider", + line_vals=["jit", "torch_ref"], + line_names=["JIT FP8 Blockwise GEMM", "Torch Ref"], + styles=[("green", "-"), ("blue", "-")], + ylabel="us", + plot_name="fp8-blockwise-scaled-mm-performance", + args={}, + ) +) +def benchmark(m, n, k, provider): + a_fp8, b_fp8, scale_a, scale_b = _make_inputs(m, n, k) + + if provider == "jit": + fn = lambda: fp8_blockwise_scaled_mm( + a_fp8, b_fp8, scale_a, scale_b, out_dtype=torch.bfloat16 + ) + elif provider == "torch_ref": + fn = lambda: _torch_ref(a_fp8, b_fp8, scale_a, scale_b) + else: + raise ValueError(f"Unknown provider: {provider}") + + return run_benchmark(fn) + + +if __name__ == "__main__": + if not is_sm120_supported(): + print( + "[skip] fp8_blockwise_scaled_mm benchmark requires SM120 with CUDA 12.8+." + ) + sys.exit(0) + benchmark.run(print_data=True) diff --git a/sgl-kernel/tests/test_fp8_blockwise_gemm.py b/test/registered/jit/test_fp8_blockwise_gemm.py similarity index 69% rename from sgl-kernel/tests/test_fp8_blockwise_gemm.py rename to test/registered/jit/test_fp8_blockwise_gemm.py index a4438de4afc4..5d097b9d28d3 100644 --- a/sgl-kernel/tests/test_fp8_blockwise_gemm.py +++ b/test/registered/jit/test_fp8_blockwise_gemm.py @@ -1,11 +1,18 @@ -import os -import random import sys from typing import Optional, Type import pytest import torch -from sgl_kernel import fp8_blockwise_scaled_mm + +from sglang.jit_kernel.fp8_blockwise_gemm import fp8_blockwise_scaled_mm +from sglang.srt.utils import is_sm120_supported +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=30, + stage="base-b", + runner_config="1-gpu-small", +) def cdiv(a: int, b: int) -> int: @@ -25,20 +32,6 @@ def baseline_scaled_mm( out_dtype: Type[torch.dtype], bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: - # We treat N-dimensional group scaling as extended numpy-style broadcasting - # in numpy simply stretches dimensions with an extent of 1 to match the - # the target shape by repeating the data along that dimension (broadcasting) - # , we extend these semantics to say if the extent of a dimension in the - # source shape is not 1 and does not match the target shape we repeat each - # element along that dimension src_shape[dim] // target_shape[dim] times - # example if we have: - # a = [[1, 2], and target_shape = (2, 4) - # [3, 4]] - # then we would expand a to: - # a = [[1, 1, 2, 2], - # [3, 3, 4, 4]] - # NOTE this function this function does not explicitly broadcast dimensions - # with an extent of 1, since this can be done implicitly by pytorch def group_broadcast(t, shape): for i, s in enumerate(shape): if t.shape[i] != s and t.shape[i] != 1: @@ -82,13 +75,16 @@ def _test_accuracy_once(M, N, K, out_dtype, device): torch.testing.assert_close(o, o1, rtol=rtol, atol=atol) -@pytest.mark.parametrize("M", [1, 3, 5, 127, 128, 512, 1024, 4096]) -@pytest.mark.parametrize("N", [128, 512, 1024, 4096, 8192, 14080]) -@pytest.mark.parametrize("K", [512, 1024, 4096, 8192, 14080, 16384]) +@pytest.mark.skipif( + not is_sm120_supported(), reason="fp8_blockwise_scaled_mm requires SM120 (>= 12.0)" +) +@pytest.mark.parametrize("M", [1, 3, 5, 32, 48, 64, 127, 128, 512, 1024, 4096]) +@pytest.mark.parametrize("N", [128, 512, 1024, 4096, 8192]) +@pytest.mark.parametrize("K", [512, 1024, 4096, 8192]) @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16]) def test_accuracy(M, N, K, out_dtype): _test_accuracy_once(M, N, K, out_dtype, "cuda") if __name__ == "__main__": - sys.exit(pytest.main([__file__])) + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/quant/test_fp8_blockwise_row_padding.py b/test/registered/quant/test_fp8_blockwise_row_padding.py deleted file mode 100644 index 43b49b5ecc09..000000000000 --- a/test/registered/quant/test_fp8_blockwise_row_padding.py +++ /dev/null @@ -1,136 +0,0 @@ -"""Unit tests for the row-padded quant path of the cutlass FP8 blockwise linear. - -`cutlass_w8a8_block_fp8_linear_with_fallback` quantizes activations into -row-aligned buffers (`sglang_per_token_group_quant_fp8_row_padded`) so the -`fp8_blockwise_scaled_mm` wrapper's per-call mat_a/scales_a padding short- -circuits. These tests pin the invariant that this is numerically identical to -the legacy unpadded path, across both row-aligned and unaligned M. -""" - -import unittest - -import torch - -from sglang.kernels.ops.quantization.fp8_kernel import ( - fp8_dtype, - per_token_group_quant_fp8, - sglang_per_token_group_quant_fp8_row_padded, -) -from sglang.srt.layers.quantization.fp8_utils import ( - _check_cutlass_block_fp8_hardware_support, - cutlass_w8a8_block_fp8_linear_with_fallback, -) -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.test_utils import CustomTestCase - -register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-large") - -_FP8_MAX = torch.finfo(fp8_dtype).max -_BLOCK = 128 -# Cover M == 1 (greedy decode), small unaligned M (speculative draft tokens), -# the 4-row alignment boundary, and a large aligned batch. -_M_VALUES = [1, 2, 3, 4, 5, 7, 13, 16, 31, 64, 256] - - -def _quant_weight_blockwise(weight_bf16: torch.Tensor, block: int = _BLOCK): - """Block-quantize a (N, K) bf16 weight to fp8 with (N//block, K//block) fp32 scales.""" - n, k = weight_bf16.shape - assert n % block == 0 and k % block == 0 - w = weight_bf16.float().reshape(n // block, block, k // block, block) - amax = w.abs().amax(dim=(1, 3)).clamp(min=1e-12) # (N//block, K//block) - scale = amax / _FP8_MAX - wq = (w / scale[:, None, :, None]).clamp(-_FP8_MAX, _FP8_MAX).to(fp8_dtype) - return wq.reshape(n, k), scale.to(torch.float32) - - -def _legacy_cutlass_linear(x_2d, weight, weight_scale): - """The pre-optimization path: unpadded quant, wrapper pads internally.""" - from sgl_kernel import fp8_blockwise_scaled_mm - - q_input, x_scale = per_token_group_quant_fp8(x_2d, _BLOCK, column_major_scales=True) - return fp8_blockwise_scaled_mm( - q_input, weight.T, x_scale, weight_scale.T, out_dtype=x_2d.dtype - ) - - -@unittest.skipUnless( - _check_cutlass_block_fp8_hardware_support(), - "cutlass block FP8 requires Hopper (SM90) or newer", -) -class TestFP8BlockwiseRowPadding(CustomTestCase): - @classmethod - def setUpClass(cls): - cls.K = 512 - cls.N = 256 - torch.manual_seed(0) - - def test_quant_buffers_row_aligned(self): - """Row-padded quant returns 4-aligned, M-major buffers whose live rows - match the legacy column-major quant bit-for-bit.""" - for m in _M_VALUES: - x = torch.randn(m, self.K, device="cuda", dtype=torch.bfloat16) * 0.1 - xq, xs = sglang_per_token_group_quant_fp8_row_padded(x, _BLOCK) - m_pad = (m + 3) // 4 * 4 - - self.assertEqual(xq.shape, (m_pad, self.K), f"M={m}") - self.assertEqual(xs.shape[0], m_pad, f"M={m}") - # scales_a must stay M-major (stride(0) == 1) for the kernel contract. - self.assertEqual(xs.stride(0), 1, f"M={m}") - - xq_ref, xs_ref = per_token_group_quant_fp8( - x, _BLOCK, column_major_scales=True - ) - self.assertEqual(xq_ref.shape, (m, self.K), f"M={m}") - # Live rows are produced by the same kernel, so they must be identical. - self.assertTrue( - torch.equal(xq[:m].view(torch.uint8), xq_ref.view(torch.uint8)), - f"quantized activation mismatch at M={m}", - ) - torch.testing.assert_close(xs[:m], xs_ref, atol=0.0, rtol=0.0) - - def test_gemm_bit_exact_vs_legacy(self): - """The full linear (row-padded) is bit-identical to the legacy unpadded GEMM.""" - weight_bf16 = ( - torch.randn(self.N, self.K, device="cuda", dtype=torch.bfloat16) * 0.1 - ) - weight, weight_scale = _quant_weight_blockwise(weight_bf16) - - for m in _M_VALUES: - x = torch.randn(m, self.K, device="cuda", dtype=torch.bfloat16) * 0.1 - - out_ref = _legacy_cutlass_linear(x, weight, weight_scale) - out_new = cutlass_w8a8_block_fp8_linear_with_fallback( - input=x, - weight=weight, - block_size=[_BLOCK, _BLOCK], - weight_scale=weight_scale, - ) - - self.assertEqual(out_new.shape, (m, self.N), f"M={m}") - self.assertTrue( - torch.equal(out_ref, out_new), - f"row-padded GEMM differs from legacy at M={m}: " - f"max_abs_diff={(out_ref.float() - out_new.float()).abs().max().item()}", - ) - - def test_linear_matches_bf16_reference(self): - """Sanity: the FP8 linear stays close to a bf16 reference matmul.""" - weight_bf16 = ( - torch.randn(self.N, self.K, device="cuda", dtype=torch.bfloat16) * 0.1 - ) - weight, weight_scale = _quant_weight_blockwise(weight_bf16) - - for m in [1, 5, 64]: - x = torch.randn(m, self.K, device="cuda", dtype=torch.bfloat16) * 0.1 - ref = (x.float() @ weight_bf16.float().T).to(torch.bfloat16) - out = cutlass_w8a8_block_fp8_linear_with_fallback( - input=x, - weight=weight, - block_size=[_BLOCK, _BLOCK], - weight_scale=weight_scale, - ) - torch.testing.assert_close(out, ref, atol=0.5, rtol=0.1) - - -if __name__ == "__main__": - unittest.main() From 27b568ccdc30b8bb974aebfed6c858451a58e15a Mon Sep 17 00:00:00 2001 From: root Date: Wed, 15 Jul 2026 07:55:08 +0000 Subject: [PATCH 11/25] fix conflict --- .../ops}/attention/dsv4/unified_kv_kernels/env_gate.py | 0 .../attention/dsv4/unified_kv_kernels/paged_decode.py | 0 .../dsv4/unified_kv_kernels/paged_decode_indices.py | 0 .../attention/dsv4/unified_kv_kernels/paged_prefill.py | 0 .../ops}/attention/dsv4/unified_kv_kernels/runtime.py | 6 +++--- python/sglang/srt/arg_groups/hisparse_hook.py | 2 +- .../layers/attention/deepseek_v4_backend_hip_radix.py | 10 +++++----- .../sglang/srt/layers/attention/dsv4/compressor_v2.py | 2 +- python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py | 4 ++-- python/sglang/srt/mem_cache/swa_radix_cache.py | 2 +- python/sglang/srt/models/deepseek_v4.py | 4 ++-- python/sglang/srt/models/deepseek_v4_dspark.py | 6 +++--- .../speculative/dspark_components/dspark_kv_inject.py | 4 ++-- .../srt/speculative/dspark_components/dspark_verify.py | 2 +- .../speculative/dspark_components/dspark_worker_v2.py | 4 ++-- 15 files changed, 23 insertions(+), 23 deletions(-) rename python/sglang/{srt/layers => kernels/ops}/attention/dsv4/unified_kv_kernels/env_gate.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/dsv4/unified_kv_kernels/paged_decode.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/dsv4/unified_kv_kernels/paged_decode_indices.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/dsv4/unified_kv_kernels/paged_prefill.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/dsv4/unified_kv_kernels/runtime.py (98%) diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/env_gate.py similarity index 100% rename from python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py rename to python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/env_gate.py diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode.py b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode.py similarity index 100% rename from python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode.py rename to python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode.py diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode_indices.py b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode_indices.py similarity index 100% rename from python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode_indices.py rename to python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode_indices.py diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_prefill.py similarity index 100% rename from python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py rename to python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_prefill.py diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/runtime.py b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py similarity index 98% rename from python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/runtime.py rename to python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py index 2c6a5677a963..bc64734987fc 100644 --- a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/runtime.py +++ b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py @@ -31,13 +31,13 @@ import triton import triton.language as tl -from sglang.srt.layers.attention.dsv4.unified_kv_kernels.paged_decode import ( +from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.paged_decode import ( sparse_attn_v4_paged_decode, ) -from sglang.srt.layers.attention.dsv4.unified_kv_kernels.paged_decode_indices import ( +from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.paged_decode_indices import ( write_v4_paged_decode_indices, ) -from sglang.srt.layers.attention.dsv4.unified_kv_kernels.paged_prefill import ( +from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.paged_prefill import ( sparse_attn_v4_paged_prefill, ) diff --git a/python/sglang/srt/arg_groups/hisparse_hook.py b/python/sglang/srt/arg_groups/hisparse_hook.py index d6c437bdbb66..c59784f6ca4c 100644 --- a/python/sglang/srt/arg_groups/hisparse_hook.py +++ b/python/sglang/srt/arg_groups/hisparse_hook.py @@ -107,7 +107,7 @@ def validate_hisparse(server_args: ServerArgs) -> None: # In unified-KV mode c4_kv_pool is None, so DeepSeekV4HiSparseTokenToKVPoolAllocator # cannot attach and pool init dies with a cryptic AssertionError. Fail fast # at startup with a clear message instead. Remove once unified-KV HiSparse lands. - from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 95bb0811f507..7bf02f6078a5 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -1125,13 +1125,13 @@ def _attach_unified_kv_decode_streams( self, core: DSV4AttnMetadata, req_pool_indices: torch.Tensor ) -> None: """build the ragged decode index streams once per forward""" - from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) if not is_unified_kv_triton(): return - from sglang.srt.layers.attention.dsv4.unified_kv_kernels import runtime + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime pool = self.token_to_kv_pool N = core.positions_casual.shape[0] @@ -1171,7 +1171,7 @@ def _attach_unified_kv_prefill_meta( seq_lens: torch.Tensor, extend_seq_lens: torch.Tensor, ) -> None: - from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) @@ -1206,7 +1206,7 @@ def _forward_unified_kv( save_kv_cache: bool = True, ) -> torch.Tensor: """unified_kv paged-attention path over the bf16 unified_kv""" - from sglang.srt.layers.attention.dsv4.unified_kv_kernels import runtime + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime pool = self.token_to_kv_pool layer_id = layer.layer_id @@ -1475,7 +1475,7 @@ def forward( token_to_kv_pool = self.token_to_kv_pool assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool) - from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index a783eac8aaff..bf682df52309 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -487,7 +487,7 @@ def forward_unified( kv_score_input = compressor.compute_kv_score(x, forward_batch) state_pool = compressor.get_state_pool(self) - from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index c75438c8010d..f2e9ffc640df 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -565,7 +565,7 @@ def __init__( c4_page_size = page_size // 4 c128_page_size = page_size // 128 - from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) @@ -1214,7 +1214,7 @@ def set_unified_key_buffer_radix_fused_norm_rope( (uncommitted verify tokens) are skipped by the scatter. """ from sglang.jit_kernel.dsv4 import fused_norm_rope_inplace - from sglang.srt.layers.attention.dsv4.unified_kv_kernels import runtime + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime fused_norm_rope_inplace(kv, kv_weight, eps, freqs_cis, positions) runtime.scatter_bf16_into_unified( diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 013a87001f6f..51651673429c 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -383,7 +383,7 @@ def swa_reprefill_tail_tokens(self) -> int: No-op (0) for the index-addressed SWA pool, whose slots are content-stable and safe to reuse. """ - from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 00377253af90..5c345cd08614 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -876,7 +876,7 @@ def _forward_prepare( use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch) kv: Optional[torch.Tensor] - from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) @@ -1130,7 +1130,7 @@ def forward( # (no DSA-CP), pass `q` as a sentinel for the `k is v` assert; the # attention path doesn't read it once `save_kv_cache=False`. attn_k = kv if kv is not None else q - from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index 9b122a4f2a24..eb92fded98ee 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -9,11 +9,11 @@ from torch import nn from sglang.jit_kernel.dsv4 import fused_q_norm_rope, fused_rope_inplace -from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config -from sglang.srt.environ import envs -from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( +from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) +from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config +from sglang.srt.environ import envs from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.quantization.base_config import QuantizationConfig diff --git a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py index 27b23e182bdf..ff0bf893dbd4 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py @@ -2,10 +2,10 @@ import torch -from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func -from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( +from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) +from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.speculative.dspark_components.kernels.dspark_verify_window import ( BuildCommitInjectLayout, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index 9173143e7406..ae40b9429719 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -5,7 +5,7 @@ import msgspec import torch -from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( +from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index abaa0f70c557..8ae96ab727c1 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -4,10 +4,10 @@ import torch -from sglang.srt.environ import envs -from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( +from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) +from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker From b2d06a7355a05ec5db684d4944f4b27ceb0ae504 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Thu, 16 Jul 2026 03:08:40 +0000 Subject: [PATCH 12/25] mirror the cuda backend about draft model --- .../deepseek_v4_backend_hip_radix.py | 161 ++++++++++++++++-- 1 file changed, 149 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 60e6237b2d74..c6b5b26c2239 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -459,12 +459,31 @@ def __init__( self.speculative_num_draft_tokens: int = ( model_runner.server_args.speculative_num_draft_tokens ) - if ( + # In-backend DSpark draft-block attention (mirrors the CUDA backend's + # is_dspark_draft path): the DSpark draft runs a fixed-gamma block verify + # that only needs SWA-window attention (the draft forces compress_ratio==0), + # so route it to a dedicated block metadata builder with need_compress=False + # to skip the c4/c128 compressor metadata the draft never reads. + # Enabled on both KV layouts: + # * non-unified (paged fp8 SWA): per-request SWA page indices via + # full_to_swa (get_dspark_swa_page_indices); + # * unified_kv (bf16 ring): _forward_unified_kv ignores swa_page_indices + # and reads the ring via _attach_unified_kv_prefill_meta, so the dspark + # gather is skipped and the generic (unused) swa_page_indices is kept. + # The block path keeps the full speculative_num_draft_tokens and derives + # gamma = value - 1 locally (graph capture) or from spec_info.draft_token_num + # (eager); since every dspark-draft verify now takes the block path, no + # global speculative_num_draft_tokens decrement is needed. + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( + is_unified_kv_triton, + ) + + self._unified_kv_triton: bool = is_unified_kv_triton() + self.is_dspark_draft: bool = ( self.speculative_num_draft_tokens is not None and getattr(model_runner, "is_draft_worker", False) and model_runner.spec_algorithm.is_dspark() - ): - self.speculative_num_draft_tokens -= 1 + ) self.speculative_step_id = speculative_step_id self.forward_metadata: Union[ DSV4Metadata, @@ -543,6 +562,7 @@ def init_forward_metadata_prefill( use_prefill_cuda_graph: bool = False, compress_gpu_plan: bool = False, extend_start_loc: Optional[torch.Tensor] = None, + dspark_block_size: Optional[int] = None, ) -> DSV4Metadata: if extend_start_loc is not None: from sglang.srt.layers.attention.dsv4.attn_metadata_kernels import ( @@ -577,6 +597,7 @@ def init_forward_metadata_prefill( out_loc=out_cache_loc, need_compress=need_compress, is_prefill=True, + dspark_block_size=dspark_block_size, ) self._attach_unified_kv_prefill_meta( core_attn_metadata, req_pool_indices, seq_lens, extend_seq_lens @@ -818,6 +839,38 @@ def init_forward_metadata_draft_extend( use_prefill_cuda_graph=use_prefill_cuda_graph, ) + def init_forward_metadata_dspark_draft_block( + self, + max_seq_len: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + out_cache_loc: torch.Tensor, + block_size: int, + ) -> DSV4Metadata: + # DSpark draft block: every request extends by a uniform gamma + # (block_size) verify tokens. Build the ragged extend layout on device + # (extend_start_loc) so the causal expansion + SWA page indices stay + # graph-safe with no host-side per-request loop. need_compress=False: + # the draft block reuses the compressed KV the target already produced. + bs = seq_lens.shape[0] + seq_lens_extended = seq_lens + block_size + extend_seq_lens = torch.full((bs,), block_size, **self.cuda_int32_kwargs) + extend_start_loc = torch.arange(bs, **self.cuda_int32_kwargs) * block_size + num_tokens = block_size * bs + return self.init_forward_metadata_prefill( + max_seq_len=max_seq_len, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens_extended, + seq_lens_cpu=None, + out_cache_loc=out_cache_loc, + num_tokens=num_tokens, + extend_seq_lens=extend_seq_lens, + extend_seq_lens_cpu=None, + need_compress=False, + extend_start_loc=extend_start_loc, + dspark_block_size=block_size, + ) + def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None: # Upgrade Raw->Full so the c4/c128 compress + core_attn + indexer # materialization is recorded inside the cuda graph; a no-op (Full @@ -933,6 +986,26 @@ def init_forward_metadata_out_graph( seq_lens=seq_lens, out_cache_loc=out_cache_loc_padded, ) + elif bucket == _GraphBucket.TARGET_VERIFY and self.is_dspark_draft: + # DSpark draft-block verify: fixed gamma per request, so the token + # count (block_size * bs) is fixed for a given bs and the graph stays + # keyed by bs (graph_key already == bs). + block_size = self.speculative_num_draft_tokens - 1 + num_tokens_block = block_size * bs + assert out_cache_loc is not None + out_cache_loc_padded = torch.nn.functional.pad( + out_cache_loc, + pad=(0, num_tokens_block - len(out_cache_loc)), + mode="constant", + value=0, + ) + temp_metadata = self.init_forward_metadata_dspark_draft_block( + max_seq_len=chosen_max_seq_len, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + out_cache_loc=out_cache_loc_padded, + block_size=block_size, + ) elif bucket == _GraphBucket.TARGET_VERIFY: assert out_cache_loc is not None ragged_layout = resolve_ragged_verify_layout(forward_batch) @@ -1027,6 +1100,17 @@ def init_forward_metadata(self, forward_batch: ForwardBatch) -> None: seq_lens=seq_lens, out_cache_loc=out_cache_loc, ) + elif self.is_dspark_draft and forward_batch.forward_mode.is_target_verify(): + # DSpark draft block: route to the in-backend block attention path + # (both KV layouts). block_size (gamma) comes from the live spec batch. + block_size = int(forward_batch.spec_info.draft_token_num) + metadata = self.init_forward_metadata_dspark_draft_block( + max_seq_len=max_seq_len, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + out_cache_loc=forward_batch.out_cache_loc, + block_size=block_size, + ) elif forward_batch.forward_mode.is_target_verify(): ragged_layout = resolve_ragged_verify_layout(forward_batch) metadata = self.init_forward_metadata_target_verify( @@ -1648,22 +1732,36 @@ def make_core_attn_metadata( out_loc: torch.Tensor, need_compress: bool = True, is_prefill: bool = False, + dspark_block_size: Optional[int] = None, ) -> DSV4AttnMetadata: assert self.swa_page_size == SWA_WINDOW seq_lens_casual = seq_lens_casual.to(torch.int32) - swa_page_indices = self.get_swa_page_indices( - seq_lens_casual=seq_lens_casual, - req_pool_indices_repeated=req_pool_indices_repeated, - ) - - swa_page_indices = _pad_last_dim( - swa_page_indices, multiples_of=PAGE_INDEX_ALIGNED_SIZE - ) + if dspark_block_size is not None and not self._unified_kv_triton: + # DSpark draft block (non-unified paged SWA): all block_size query + # tokens of a request share one SWA window (based on the block's + # prefix), so build the page indices per-request (bs rows) instead of + # the generic per-token [num_q, SWA_WINDOW] gather. Under unified_kv + # _forward_unified_kv ignores swa_page_indices (it reads the bf16 ring + # via pf_* metadata), so fall through to the generic path there. + swa_page_indices, swa_topk_lengths = self.get_dspark_swa_page_indices( + seq_lens_casual=seq_lens_casual, + req_pool_indices_repeated=req_pool_indices_repeated, + out_loc=out_loc, + block_size=dspark_block_size, + ) + else: + swa_page_indices = self.get_swa_page_indices( + seq_lens_casual=seq_lens_casual, + req_pool_indices_repeated=req_pool_indices_repeated, + ) + swa_page_indices = _pad_last_dim( + swa_page_indices, multiples_of=PAGE_INDEX_ALIGNED_SIZE + ) + swa_topk_lengths = torch.clamp(seq_lens_casual, max=SWA_WINDOW) raw_positions = seq_lens_casual - 1 - swa_topk_lengths = torch.clamp(seq_lens_casual, max=SWA_WINDOW) page_table = req_to_token[ req_pool_indices_repeated, : max_seq_len : self.page_size @@ -1716,6 +1814,45 @@ def get_swa_page_indices( # flash_mla attention requires int32 page indices. return swa_indices.to(torch.int32) + def get_dspark_swa_page_indices( + self, + *, + seq_lens_casual: torch.Tensor, + req_pool_indices_repeated: torch.Tensor, + out_loc: torch.Tensor, + block_size: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Mirrors the CUDA backend: all block_size query tokens of a request + # share one SWA window (based on the block's prefix), so gather per + # request (bs rows) rather than per token. Reuses the shared triton/torch + # DSpark window kernels (dispatched by input placement; the triton path + # runs on ROCm). Only valid on the non-unified paged-SWA KV layout, where + # full_to_swa_index_mapping exists. + from sglang.srt.speculative.dspark_components.kernels.dspark_attn_metadata import ( + BuildDsparkSwaPageIndices, + ComputeDsparkWindowGather, + ) + + gather = ComputeDsparkWindowGather.execute( + seq_lens_casual=seq_lens_casual, + req_pool_indices_repeated=req_pool_indices_repeated, + block_size=block_size, + swa_window=SWA_WINDOW, + ) + swa_page_indices, swa_topk_lengths = BuildDsparkSwaPageIndices.execute( + req_to_token=self.req_to_token, + full_to_swa_mapping=self.token_to_kv_pool.full_to_swa_index_mapping, + req_pool_indices_per_request=gather.req_pool_indices_per_request, + offsets=gather.offsets, + invalid=gather.invalid, + out_loc=out_loc[: gather.num_q], + context_lens=gather.context_lens, + block_size=block_size, + swa_window=SWA_WINDOW, + page_index_aligned_size=PAGE_INDEX_ALIGNED_SIZE, + ) + return swa_page_indices, swa_topk_lengths + class DeepseekV4MultiStepBackend(DeepseekV4HipRadixBackend): def __init__( From e90a20ac07df6e52d62a80eb8841bc9a2433bc8f Mon Sep 17 00:00:00 2001 From: At1a8 Date: Thu, 16 Jul 2026 05:46:20 +0000 Subject: [PATCH 13/25] Revert "mirror the cuda backend about draft model" This reverts commit b2d06a7355a05ec5db684d4944f4b27ceb0ae504. --- .../deepseek_v4_backend_hip_radix.py | 161 ++---------------- 1 file changed, 12 insertions(+), 149 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index c6b5b26c2239..60e6237b2d74 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -459,31 +459,12 @@ def __init__( self.speculative_num_draft_tokens: int = ( model_runner.server_args.speculative_num_draft_tokens ) - # In-backend DSpark draft-block attention (mirrors the CUDA backend's - # is_dspark_draft path): the DSpark draft runs a fixed-gamma block verify - # that only needs SWA-window attention (the draft forces compress_ratio==0), - # so route it to a dedicated block metadata builder with need_compress=False - # to skip the c4/c128 compressor metadata the draft never reads. - # Enabled on both KV layouts: - # * non-unified (paged fp8 SWA): per-request SWA page indices via - # full_to_swa (get_dspark_swa_page_indices); - # * unified_kv (bf16 ring): _forward_unified_kv ignores swa_page_indices - # and reads the ring via _attach_unified_kv_prefill_meta, so the dspark - # gather is skipped and the generic (unused) swa_page_indices is kept. - # The block path keeps the full speculative_num_draft_tokens and derives - # gamma = value - 1 locally (graph capture) or from spec_info.draft_token_num - # (eager); since every dspark-draft verify now takes the block path, no - # global speculative_num_draft_tokens decrement is needed. - from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( - is_unified_kv_triton, - ) - - self._unified_kv_triton: bool = is_unified_kv_triton() - self.is_dspark_draft: bool = ( + if ( self.speculative_num_draft_tokens is not None and getattr(model_runner, "is_draft_worker", False) and model_runner.spec_algorithm.is_dspark() - ) + ): + self.speculative_num_draft_tokens -= 1 self.speculative_step_id = speculative_step_id self.forward_metadata: Union[ DSV4Metadata, @@ -562,7 +543,6 @@ def init_forward_metadata_prefill( use_prefill_cuda_graph: bool = False, compress_gpu_plan: bool = False, extend_start_loc: Optional[torch.Tensor] = None, - dspark_block_size: Optional[int] = None, ) -> DSV4Metadata: if extend_start_loc is not None: from sglang.srt.layers.attention.dsv4.attn_metadata_kernels import ( @@ -597,7 +577,6 @@ def init_forward_metadata_prefill( out_loc=out_cache_loc, need_compress=need_compress, is_prefill=True, - dspark_block_size=dspark_block_size, ) self._attach_unified_kv_prefill_meta( core_attn_metadata, req_pool_indices, seq_lens, extend_seq_lens @@ -839,38 +818,6 @@ def init_forward_metadata_draft_extend( use_prefill_cuda_graph=use_prefill_cuda_graph, ) - def init_forward_metadata_dspark_draft_block( - self, - max_seq_len: int, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - out_cache_loc: torch.Tensor, - block_size: int, - ) -> DSV4Metadata: - # DSpark draft block: every request extends by a uniform gamma - # (block_size) verify tokens. Build the ragged extend layout on device - # (extend_start_loc) so the causal expansion + SWA page indices stay - # graph-safe with no host-side per-request loop. need_compress=False: - # the draft block reuses the compressed KV the target already produced. - bs = seq_lens.shape[0] - seq_lens_extended = seq_lens + block_size - extend_seq_lens = torch.full((bs,), block_size, **self.cuda_int32_kwargs) - extend_start_loc = torch.arange(bs, **self.cuda_int32_kwargs) * block_size - num_tokens = block_size * bs - return self.init_forward_metadata_prefill( - max_seq_len=max_seq_len, - req_pool_indices=req_pool_indices, - seq_lens=seq_lens_extended, - seq_lens_cpu=None, - out_cache_loc=out_cache_loc, - num_tokens=num_tokens, - extend_seq_lens=extend_seq_lens, - extend_seq_lens_cpu=None, - need_compress=False, - extend_start_loc=extend_start_loc, - dspark_block_size=block_size, - ) - def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None: # Upgrade Raw->Full so the c4/c128 compress + core_attn + indexer # materialization is recorded inside the cuda graph; a no-op (Full @@ -986,26 +933,6 @@ def init_forward_metadata_out_graph( seq_lens=seq_lens, out_cache_loc=out_cache_loc_padded, ) - elif bucket == _GraphBucket.TARGET_VERIFY and self.is_dspark_draft: - # DSpark draft-block verify: fixed gamma per request, so the token - # count (block_size * bs) is fixed for a given bs and the graph stays - # keyed by bs (graph_key already == bs). - block_size = self.speculative_num_draft_tokens - 1 - num_tokens_block = block_size * bs - assert out_cache_loc is not None - out_cache_loc_padded = torch.nn.functional.pad( - out_cache_loc, - pad=(0, num_tokens_block - len(out_cache_loc)), - mode="constant", - value=0, - ) - temp_metadata = self.init_forward_metadata_dspark_draft_block( - max_seq_len=chosen_max_seq_len, - req_pool_indices=req_pool_indices, - seq_lens=seq_lens, - out_cache_loc=out_cache_loc_padded, - block_size=block_size, - ) elif bucket == _GraphBucket.TARGET_VERIFY: assert out_cache_loc is not None ragged_layout = resolve_ragged_verify_layout(forward_batch) @@ -1100,17 +1027,6 @@ def init_forward_metadata(self, forward_batch: ForwardBatch) -> None: seq_lens=seq_lens, out_cache_loc=out_cache_loc, ) - elif self.is_dspark_draft and forward_batch.forward_mode.is_target_verify(): - # DSpark draft block: route to the in-backend block attention path - # (both KV layouts). block_size (gamma) comes from the live spec batch. - block_size = int(forward_batch.spec_info.draft_token_num) - metadata = self.init_forward_metadata_dspark_draft_block( - max_seq_len=max_seq_len, - req_pool_indices=req_pool_indices, - seq_lens=seq_lens, - out_cache_loc=forward_batch.out_cache_loc, - block_size=block_size, - ) elif forward_batch.forward_mode.is_target_verify(): ragged_layout = resolve_ragged_verify_layout(forward_batch) metadata = self.init_forward_metadata_target_verify( @@ -1732,36 +1648,22 @@ def make_core_attn_metadata( out_loc: torch.Tensor, need_compress: bool = True, is_prefill: bool = False, - dspark_block_size: Optional[int] = None, ) -> DSV4AttnMetadata: assert self.swa_page_size == SWA_WINDOW seq_lens_casual = seq_lens_casual.to(torch.int32) - if dspark_block_size is not None and not self._unified_kv_triton: - # DSpark draft block (non-unified paged SWA): all block_size query - # tokens of a request share one SWA window (based on the block's - # prefix), so build the page indices per-request (bs rows) instead of - # the generic per-token [num_q, SWA_WINDOW] gather. Under unified_kv - # _forward_unified_kv ignores swa_page_indices (it reads the bf16 ring - # via pf_* metadata), so fall through to the generic path there. - swa_page_indices, swa_topk_lengths = self.get_dspark_swa_page_indices( - seq_lens_casual=seq_lens_casual, - req_pool_indices_repeated=req_pool_indices_repeated, - out_loc=out_loc, - block_size=dspark_block_size, - ) - else: - swa_page_indices = self.get_swa_page_indices( - seq_lens_casual=seq_lens_casual, - req_pool_indices_repeated=req_pool_indices_repeated, - ) - swa_page_indices = _pad_last_dim( - swa_page_indices, multiples_of=PAGE_INDEX_ALIGNED_SIZE - ) - swa_topk_lengths = torch.clamp(seq_lens_casual, max=SWA_WINDOW) + swa_page_indices = self.get_swa_page_indices( + seq_lens_casual=seq_lens_casual, + req_pool_indices_repeated=req_pool_indices_repeated, + ) + + swa_page_indices = _pad_last_dim( + swa_page_indices, multiples_of=PAGE_INDEX_ALIGNED_SIZE + ) raw_positions = seq_lens_casual - 1 + swa_topk_lengths = torch.clamp(seq_lens_casual, max=SWA_WINDOW) page_table = req_to_token[ req_pool_indices_repeated, : max_seq_len : self.page_size @@ -1814,45 +1716,6 @@ def get_swa_page_indices( # flash_mla attention requires int32 page indices. return swa_indices.to(torch.int32) - def get_dspark_swa_page_indices( - self, - *, - seq_lens_casual: torch.Tensor, - req_pool_indices_repeated: torch.Tensor, - out_loc: torch.Tensor, - block_size: int, - ) -> Tuple[torch.Tensor, torch.Tensor]: - # Mirrors the CUDA backend: all block_size query tokens of a request - # share one SWA window (based on the block's prefix), so gather per - # request (bs rows) rather than per token. Reuses the shared triton/torch - # DSpark window kernels (dispatched by input placement; the triton path - # runs on ROCm). Only valid on the non-unified paged-SWA KV layout, where - # full_to_swa_index_mapping exists. - from sglang.srt.speculative.dspark_components.kernels.dspark_attn_metadata import ( - BuildDsparkSwaPageIndices, - ComputeDsparkWindowGather, - ) - - gather = ComputeDsparkWindowGather.execute( - seq_lens_casual=seq_lens_casual, - req_pool_indices_repeated=req_pool_indices_repeated, - block_size=block_size, - swa_window=SWA_WINDOW, - ) - swa_page_indices, swa_topk_lengths = BuildDsparkSwaPageIndices.execute( - req_to_token=self.req_to_token, - full_to_swa_mapping=self.token_to_kv_pool.full_to_swa_index_mapping, - req_pool_indices_per_request=gather.req_pool_indices_per_request, - offsets=gather.offsets, - invalid=gather.invalid, - out_loc=out_loc[: gather.num_q], - context_lens=gather.context_lens, - block_size=block_size, - swa_window=SWA_WINDOW, - page_index_aligned_size=PAGE_INDEX_ALIGNED_SIZE, - ) - return swa_page_indices, swa_topk_lengths - class DeepseekV4MultiStepBackend(DeepseekV4HipRadixBackend): def __init__( From 0227778bd3e2aa9d19d2cd460f95c53dea6174ff Mon Sep 17 00:00:00 2001 From: At1a8 Date: Fri, 17 Jul 2026 06:31:18 +0000 Subject: [PATCH 14/25] fix kk comments --- .../dsv4/unified_kv_kernels/runtime.py | 1 + .../deepseek_v4_backend_hip_radix.py | 30 ++++++++++++------- .../dspark_components/dspark_verify.py | 1 - 3 files changed, 20 insertions(+), 12 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py index bc64734987fc..79ae94ab03dd 100644 --- a/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py +++ b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py @@ -149,6 +149,7 @@ def scatter_bf16_into_unified( return assert kv.is_contiguous() and kv.dtype == unified_kv.dtype assert loc.is_contiguous() + assert unified_kv.is_contiguous() _scatter_loc_kernel[(n_rows,)]( kv, loc, diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 60e6237b2d74..a94ae1dd5823 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -459,12 +459,18 @@ def __init__( self.speculative_num_draft_tokens: int = ( model_runner.server_args.speculative_num_draft_tokens ) - if ( - self.speculative_num_draft_tokens is not None - and getattr(model_runner, "is_draft_worker", False) + self.is_dspark_draft = ( + getattr(model_runner, "is_draft_worker", False) and model_runner.spec_algorithm.is_dspark() - ): - self.speculative_num_draft_tokens -= 1 + ) + self.target_verify_num_draft_tokens = self.speculative_num_draft_tokens + if self.is_dspark_draft: + assert self.speculative_num_draft_tokens is not None + assert self.speculative_num_draft_tokens > 1 + # DSpark draft workers verify gamma rows. The server arg keeps the + # CUDA-side convention gamma + 1, so use an explicit effective value + # instead of mutating speculative_num_draft_tokens in place. + self.target_verify_num_draft_tokens = self.speculative_num_draft_tokens - 1 self.speculative_step_id = speculative_step_id self.forward_metadata: Union[ DSV4Metadata, @@ -681,10 +687,12 @@ def init_forward_metadata_target_verify_old( extend_seq_lens_cpu = None seq_lens_cpu = None else: - seq_lens = seq_lens + self.speculative_num_draft_tokens - seq_lens_cpu = [x + self.speculative_num_draft_tokens for x in seq_lens_cpu] - extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * batch_size - num_tokens = self.speculative_num_draft_tokens * batch_size + seq_lens = seq_lens + self.target_verify_num_draft_tokens + seq_lens_cpu = [ + x + self.target_verify_num_draft_tokens for x in seq_lens_cpu + ] + extend_seq_lens_cpu = [self.target_verify_num_draft_tokens] * batch_size + num_tokens = self.target_verify_num_draft_tokens * batch_size extend_seq_lens = self._move_to_device(extend_seq_lens_cpu) if out_cache_loc is None: out_cache_loc = seq_lens.new_zeros(num_tokens) @@ -710,7 +718,7 @@ def make_forward_metadata_from_raw_verify( seq_lens = raw_metadata.seq_lens out_cache_loc = raw_metadata.out_cache_loc - bs, num_draft_tokens = len(seq_lens), self.speculative_num_draft_tokens + bs, num_draft_tokens = len(seq_lens), self.target_verify_num_draft_tokens seq_lens = seq_lens + num_draft_tokens extend_seq_lens = raw_metadata.extend_seq_lens if extend_seq_lens is None or extend_seq_lens.numel() != bs: @@ -941,7 +949,7 @@ def init_forward_metadata_out_graph( num_tokens_v = ragged_layout.graph_num_tokens graph_key = num_tokens_v else: - num_tokens_v = self.speculative_num_draft_tokens * bs + num_tokens_v = self.target_verify_num_draft_tokens * bs out_cache_loc_padded = torch.nn.functional.pad( out_cache_loc, pad=(0, num_tokens_v - len(out_cache_loc)), diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index ae40b9429719..c8f4c2f567cc 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -309,7 +309,6 @@ def commit_hidden( raise RuntimeError("DSpark verify requires target hidden states, got None.") hidden = hidden.view(bs, self.verify_num_draft_tokens, -1) state_slot = None - pool = self.kv_injector.draft_model_runner.token_to_kv_pool if is_unified_kv_triton(): # unified_kv needs the per-token draft req slot to address the SWA ring # (state_slot * ring + pos % ring). Verify tokens are the latest in each From 7a9fd6a68110fe5d81ab1568672cd57380288a58 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Fri, 17 Jul 2026 08:44:09 +0000 Subject: [PATCH 15/25] ci test for kernel --- .../amd/test_deepseek_v4_pro_fp4_dspark.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py b/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py index 7fda67bda8e2..45959cc2d375 100644 --- a/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py +++ b/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py @@ -13,7 +13,10 @@ from types import SimpleNamespace import requests +import torch +from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime +from sglang.srt.speculative.dspark_components.kernels import dspark_verify_window from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_amd_ci from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k @@ -35,6 +38,57 @@ SERVER_LAUNCH_TIMEOUT = 5400 GSM8K_ACCURACY_THRESHOLD = 0.92 AVG_SPEC_ACCEPT_LENGTH_THRESHOLD = 3.0 +DEVICE = torch.device("cuda") + + +class TestDSparkUnifiedKVKernelsAMD(CustomTestCase): + def test_build_unified_commit_inject_layout(self): + stride, ring_stride = 7, 128 + req_pool_indices = torch.tensor([3, 0, 5, 1], device=DEVICE, dtype=torch.int32) + prefix_lens = torch.tensor( + [10, 127, 128, 255], device=DEVICE, dtype=torch.int64 + ) + block_pos_offsets = torch.arange(stride, device=DEVICE, dtype=torch.int64) + commit_lens = torch.tensor([0, 3, stride, 5], device=DEVICE, dtype=torch.int32) + + got = dspark_verify_window.build_unified_commit_inject_layout( + req_pool_indices=req_pool_indices, + prefix_lens=prefix_lens, + block_pos_offsets=block_pos_offsets, + commit_lens=commit_lens, + stride=stride, + ring_stride=ring_stride, + ) + + positions_2d = prefix_lens.view(-1, 1) + block_pos_offsets[:stride] + loc_2d = req_pool_indices.to(torch.int64).view(-1, 1) * ring_stride + loc_2d = loc_2d + positions_2d % ring_stride + col = torch.arange(stride, device=DEVICE).view(1, -1) + committed = col < commit_lens.to(torch.long).view(-1, 1) + ref_loc = torch.where(committed, loc_2d, torch.full_like(loc_2d, -1)).to( + torch.int32 + ) + + self.assertTrue(torch.equal(got.positions, positions_2d.reshape(-1))) + self.assertTrue(torch.equal(got.swa_loc, ref_loc.reshape(-1))) + + def test_scatter_bf16_into_unified(self): + torch.manual_seed(20) + n_rows, dim, n_pages = 8, 16, 32 + kv = torch.randn(n_rows, dim, device=DEVICE).to(torch.bfloat16).contiguous() + loc = torch.tensor( + [3, -1, 5, 7, 0, -1, 9, 11], device=DEVICE, dtype=torch.int32 + ) + unified = torch.zeros(n_pages, dim, device=DEVICE, dtype=torch.bfloat16) + expected = unified.clone() + keep = loc >= 0 + expected[loc[keep].long()] = kv[keep] + + runtime.scatter_bf16_into_unified(kv=kv, loc=loc, unified_kv=unified) + self.assertTrue(torch.equal(unified, expected)) + + with self.assertRaises(AssertionError): + runtime.scatter_bf16_into_unified(kv=kv, loc=loc, unified_kv=unified.t()) class TestDeepseekV4DSparkUnifiedKVGSM8K(CustomTestCase): From 63479f6a71270d9530b84aad1860b0aa8227cbea Mon Sep 17 00:00:00 2001 From: At1a8 Date: Fri, 17 Jul 2026 09:38:49 +0000 Subject: [PATCH 16/25] fix bug --- .../sglang/srt/speculative/dspark_components/dspark_draft.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index c4cb2e013ffb..9ab00ff3aaa8 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -415,6 +415,9 @@ def _fill_dp_moe_sync_metadata( batch.global_num_tokens_for_logprob, ) device = self.draft_model_runner.device + # Keep the raw per-rank request counts for CUDA graph batch-size checks. + # The sync tensors below are scaled for the speculative draft width. + forward_batch.original_global_num_tokens_cpu = batch.global_num_tokens forward_batch.global_num_tokens_cpu = gnt forward_batch.global_num_tokens_for_logprob_cpu = gnt_logprob forward_batch.global_num_tokens_gpu = torch.tensor(gnt, dtype=torch.int64).to( From 5c810a115292d857af730757dba3cf3f18f7a335 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Sat, 18 Jul 2026 17:08:16 +0000 Subject: [PATCH 17/25] update ci --- .../amd/test_deepseek_v4_pro_fp4_dspark.py | 52 ++++++++----------- 1 file changed, 21 insertions(+), 31 deletions(-) diff --git a/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py b/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py index 45959cc2d375..6f10ae0832c8 100644 --- a/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py +++ b/test/registered/amd/test_deepseek_v4_pro_fp4_dspark.py @@ -36,10 +36,28 @@ "DEEPSEEK_V4_DSPARK_MODEL_PATH", "deepseek-ai/DeepSeek-V4-Pro-DSpark" ) SERVER_LAUNCH_TIMEOUT = 5400 +FLASHMLA_BACKEND = os.environ.get("SGLANG_HACK_FLASHMLA_BACKEND", "unified_kv_triton") GSM8K_ACCURACY_THRESHOLD = 0.92 AVG_SPEC_ACCEPT_LENGTH_THRESHOLD = 3.0 DEVICE = torch.device("cuda") +COMMON_ENV_VARS = { + "SGLANG_DEFAULT_THINKING": "1", + "SGLANG_DSV4_REASONING_EFFORT": "max", + "SGLANG_USE_ROCM700A": "0", + "SGLANG_HACK_FLASHMLA_BACKEND": FLASHMLA_BACKEND, + "AITER_BF16_FP8_MOE_BOUND": "0", +} + +DSPARK_ENV_VARS = { + "SGLANG_RAGGED_VERIFY_MODE": "static", +} + +# FP4 variant (matches test_deepseek_v4_pro_fp4.py; V4-Pro also auto-detects it). +FP4_ENV_VARS = { + "SGLANG_DSV4_FP4_EXPERTS": "true", +} + class TestDSparkUnifiedKVKernelsAMD(CustomTestCase): def test_build_unified_commit_inject_layout(self): @@ -97,37 +115,9 @@ def setUpClass(cls): cls.model = DEEPSEEK_V4_DSPARK_MODEL_PATH cls.base_url = DEFAULT_URL_FOR_TEST env = os.environ.copy() - env.update( - { - "SGLANG_DEFAULT_THINKING": "1", - "SGLANG_DSV4_REASONING_EFFORT": "max", - "SGLANG_OPT_DEEPGEMM_HC_PRENORM": "false", - "SGLANG_USE_AITER": "1", - "SGLANG_USE_ROCM700A": "0", - "SGLANG_OPT_USE_FUSED_COMPRESS": "true", - "SGLANG_HACK_FLASHMLA_BACKEND": "unified_kv_triton", - "SGLANG_RAGGED_VERIFY_MODE": "static", - "SGLANG_DSPARK_ENABLE_SPS_ONLINE_PROFILE": "0", - "SGLANG_OPT_FP8_WO_A_GEMM": "false", - "SGLANG_OPT_USE_JIT_INDEXER_METADATA": "false", - "SGLANG_OPT_USE_TOPK_V2": "false", - "SGLANG_OPT_USE_AITER_INDEXER": "true", - "SGLANG_OPT_USE_TILELANG_INDEXER": "false", - "SGLANG_OPT_USE_TILELANG_MHC_PRE": "false", - "SGLANG_OPT_USE_TILELANG_MHC_POST": "false", - "SGLANG_FP8_PAGED_MQA_LOGITS_TORCH": "1", - "SGLANG_OPT_USE_FUSED_COMPRESS_TRITON": "true", - "SGLANG_OPT_USE_MULTI_STREAM_OVERLAP": "false", - "SGLANG_ROCM_USE_MULTI_STREAM": "false", - "AITER_BF16_FP8_MOE_BOUND": "0", - "SGLANG_EAGER_INPUT_NO_COPY": "true", - "SGLANG_SHARED_EXPERT_TP1": "1", - "SGLANG_DP_SHARED_EXPERT_LOCAL": "1", - "SGLANG_DP_USE_GATHERV": "1", - "SGLANG_DP_USE_REDUCE_SCATTER": "1", - "GPU_MAX_HW_QUEUES": "5", - } - ) + env.update(COMMON_ENV_VARS) + env.update(DSPARK_ENV_VARS) + env.update(FP4_ENV_VARS) other_args = [ "--trust-remote-code", "--tp", From a5a8e16cc0fea65e4820d3445f1485bb6b585d4a Mon Sep 17 00:00:00 2001 From: At1a8 Date: Wed, 22 Jul 2026 16:44:12 +0000 Subject: [PATCH 18/25] fix --- .../srt/layers/attention/deepseek_v4_backend_hip_radix.py | 2 +- python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py | 2 +- python/sglang/srt/models/deepseek_v4_dspark.py | 3 +-- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index a94ae1dd5823..a37f35d52a36 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -551,7 +551,7 @@ def init_forward_metadata_prefill( extend_start_loc: Optional[torch.Tensor] = None, ) -> DSV4Metadata: if extend_start_loc is not None: - from sglang.srt.layers.attention.dsv4.attn_metadata_kernels import ( + from sglang.kernels.ops.attention.dsv4_attn_metadata_kernels import ( ExpandPrefillCausally, ) diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 31d320c64c33..d77077e7253e 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -1213,7 +1213,7 @@ def set_unified_key_buffer_radix_fused_norm_rope( scatter it into ``unified_kv[swa_loc]``. Rows with swa_loc < 0 (uncommitted verify tokens) are skipped by the scatter. """ - from sglang.jit_kernel.dsv4 import fused_norm_rope_inplace + from sglang.kernels.ops.attention.dsv4 import fused_norm_rope_inplace from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime fused_norm_rope_inplace(kv, kv_weight, eps, freqs_cis, positions) diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index 59aaedafaa81..7279baa1af76 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -8,11 +8,10 @@ import torch.nn.functional as F from torch import nn -from sglang.jit_kernel.dsv4 import fused_q_norm_rope, fused_rope_inplace +from sglang.kernels.ops.attention.dsv4 import fused_q_norm_rope, fused_rope_inplace from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) -from sglang.kernels.ops.attention.dsv4 import fused_q_norm_rope, fused_rope_inplace from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config from sglang.srt.environ import envs from sglang.srt.layers.layernorm import RMSNorm From 95e824344cbce2bad27151ea8c984d61f6736330 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 31 Jul 2026 02:45:17 +0000 Subject: [PATCH 19/25] fix lint --- python/sglang/srt/speculative/dflash_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 546f9c5207f2..e4dcb25811de 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -12,8 +12,8 @@ from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod from sglang.srt.layers.sampler import apply_custom_logit_processor from sglang.srt.managers.schedule_batch import Req -from sglang.srt.utils import is_cuda, is_hip, is_musa from sglang.srt.speculative.spec_utils import _sample_simulated_acc_len +from sglang.srt.utils import is_cuda, is_hip, is_musa DEFAULT_DFLASH_MASK_TOKEN = "<|MASK|>" From 058c424d931c0c1f2135c5c3c0229ef97b6f55fe Mon Sep 17 00:00:00 2001 From: At1a8 Date: Fri, 31 Jul 2026 07:04:56 +0000 Subject: [PATCH 20/25] fix bug --- .../sglang/kernels/ops/speculative/dspark/dspark_draft_model.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py b/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py index 5915c090b706..3764cdbd0bdd 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py @@ -107,6 +107,7 @@ def _online_partial_kernel( temperature = tl.load(temperatures_ptr + row) s = logits / temperature tile_max = tl.max(s, axis=0) + tile_max = tl.where(tile_max == float("-inf"), 0.0, tile_max) greedy = tl.load(greedy_mask_ptr + row) != 0 noise = tl.load(exp_noise_ptr + row * V + offs, mask=mask, other=1.0) denom = tl.where(greedy, 1.0, noise) @@ -139,6 +140,7 @@ def _online_combine_kernel( partial_idx_ptr + row * n_tiles + offs, mask=mask, other=_IDX_SENTINEL ) global_max = tl.max(tile_max, axis=0) + global_max = tl.where(global_max == float("-inf"), 0.0, global_max) rescaled = keys * tl.exp(tile_max - global_max) rescaled = tl.where(mask, rescaled, -1.0) best = tl.max(rescaled, axis=0) From 172451ec95f71f9e64ac3bbf05c4f61c4365bf37 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Fri, 31 Jul 2026 09:40:12 +0000 Subject: [PATCH 21/25] fix --- .../sglang/kernels/ops/speculative/dspark/dspark_draft_model.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py b/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py index 3764cdbd0bdd..5915c090b706 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py @@ -107,7 +107,6 @@ def _online_partial_kernel( temperature = tl.load(temperatures_ptr + row) s = logits / temperature tile_max = tl.max(s, axis=0) - tile_max = tl.where(tile_max == float("-inf"), 0.0, tile_max) greedy = tl.load(greedy_mask_ptr + row) != 0 noise = tl.load(exp_noise_ptr + row * V + offs, mask=mask, other=1.0) denom = tl.where(greedy, 1.0, noise) @@ -140,7 +139,6 @@ def _online_combine_kernel( partial_idx_ptr + row * n_tiles + offs, mask=mask, other=_IDX_SENTINEL ) global_max = tl.max(tile_max, axis=0) - global_max = tl.where(global_max == float("-inf"), 0.0, global_max) rescaled = keys * tl.exp(tile_max - global_max) rescaled = tl.where(mask, rescaled, -1.0) best = tl.max(rescaled, axis=0) From 6f406f1f5fea82d205158709b8cb1ce4f06ed00b Mon Sep 17 00:00:00 2001 From: At1a8 Date: Fri, 31 Jul 2026 09:49:48 +0000 Subject: [PATCH 22/25] make sure hip only changes --- .../sglang/srt/speculative/dspark_components/dspark_draft.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index 2a0262429630..e1d94cba3ff5 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -26,6 +26,7 @@ spec_scale_global_num_tokens, ) from sglang.srt.speculative.spec_utils import draft_tp_context +from sglang.srt.utils import is_hip from sglang.srt.utils.invariants import Bucket, Invariant, NotNaN, expect logger = logging.getLogger(__name__) @@ -438,7 +439,8 @@ def _fill_dp_moe_sync_metadata( device = self.draft_model_runner.device # Keep the raw per-rank request counts for CUDA graph batch-size checks. # The sync tensors below are scaled for the speculative draft width. - forward_batch.original_global_num_tokens_cpu = batch.global_num_tokens + if is_hip(): + forward_batch.original_global_num_tokens_cpu = batch.global_num_tokens forward_batch.global_num_tokens_cpu = gnt forward_batch.global_num_tokens_for_logprob_cpu = gnt_logprob forward_batch.global_num_tokens_gpu = torch.tensor(gnt, dtype=torch.int64).to( From 2bbcec79711e8ee8c365ccca51ea23d99ebac428 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Sat, 1 Aug 2026 00:59:12 +0000 Subject: [PATCH 23/25] fix --- .../srt/layers/attention/deepseek_v4_backend_hip_radix.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index f52b641a9c06..f73b79aaacf9 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -456,6 +456,7 @@ def __init__( assert self.topk in [0, 1], "MTP Topk > 1 not supported for DeepSeek V4" self.mtp_enabled = self.topk > 0 self.speculative_num_steps = speculative_num_steps + self.speculative_num_draft_tokens: int = get_spec().speculative_num_draft_tokens self.is_dspark_draft = ( getattr(model_runner, "is_draft_worker", False) and model_runner.spec_algorithm.is_dspark() @@ -468,7 +469,6 @@ def __init__( # CUDA-side convention gamma + 1, so use an explicit effective value # instead of mutating speculative_num_draft_tokens in place. self.target_verify_num_draft_tokens = self.speculative_num_draft_tokens - 1 - self.speculative_num_draft_tokens: int = get_spec().speculative_num_draft_tokens self.speculative_step_id = speculative_step_id self.forward_metadata: Union[ DSV4Metadata, From ed2752fc63a34ba0f6f56f2647f3fd4a79709a35 Mon Sep 17 00:00:00 2001 From: At1a8 Date: Tue, 4 Aug 2026 08:05:42 +0000 Subject: [PATCH 24/25] x --- .../speculative/dspark_components/dspark_draft.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index 23b0c0c7f428..c7f27428ef17 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -27,7 +27,6 @@ spec_scale_global_num_tokens, ) from sglang.srt.speculative.spec_utils import draft_tp_context -from sglang.srt.utils import is_hip from sglang.srt.utils.invariants import Bucket, Invariant, NotNaN, expect logger = logging.getLogger(__name__) @@ -386,10 +385,13 @@ def _fill_dp_moe_sync_metadata( batch.global_num_tokens_for_logprob, ) device = self.draft_model_runner.device - # Keep the raw per-rank request counts for CUDA graph batch-size checks. - # The sync tensors below are scaled for the speculative draft width. - if is_hip(): - forward_batch.original_global_num_tokens_cpu = batch.global_num_tokens + forward_batch.original_global_num_tokens_cpu = batch.global_num_tokens + num_tokens = forward_batch.input_ids.numel() + if enable_num_token_non_padded(): + forward_batch.num_token_non_padded = torch.tensor( + num_tokens, dtype=torch.int32, device=device + ) + forward_batch.num_token_non_padded_cpu = num_tokens forward_batch.global_num_tokens_cpu = gnt forward_batch.global_num_tokens_for_logprob_cpu = gnt_logprob forward_batch.global_num_tokens_gpu = torch.tensor(gnt, dtype=torch.int64).to( From 7c207d84be134aa069d1609a06219c45255618cf Mon Sep 17 00:00:00 2001 From: At1a8 Date: Thu, 6 Aug 2026 07:35:54 +0000 Subject: [PATCH 25/25] lint --- python/sglang/srt/speculative/dflash_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index f5aa6e158922..649f0dd9a849 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -53,6 +53,7 @@ from sglang.kernels.ops.sampling.renorm_triton import ( top_p_renorm_probs_triton as top_p_renorm_prob, ) + tree_speculative_sampling_target_only = None else: top_k_renorm_prob = None