From b613aba943624a9005eee163bbc4084a25939fba Mon Sep 17 00:00:00 2001 From: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Date: Tue, 22 Sep 2026 05:18:22 -0700 Subject: [PATCH] kda: zero each padded CUDA-graph row's own dense output interval in the CuTe DSL MTP kernel (#138) * kda: zero each padded CUDA-graph row's own dense output interval in the CuTe DSL MTP kernel Graph-padded requests (ssm_state_indices == -1) all share the last real cu_seqlens endpoint: production graph metadata collapses every padded request's query_start_loc to the real-token endpoint (hybrid_linear_attn_backend.py, verify branch with num_padding > 0), so the kernel's pad branch cleared only the first padded request's T_LOOP rows and left the remaining padded rows as torch.empty bytes. The wrapper enforces a fixed dense 1 + num_spec width per request, so the pad branch now clears row block i_n * T_LOOP. The defect geometry is measured, not just inferred: an instrumented engine run observed the first padded request's dense output interval zero in 100% of 5.3M observations across 8 ranks while padded requests 2 and 3 were nonzero in 6.1%-59.0% of theirs, with zero nonfinite values in every arm. That experiment applied no fix. Adds test_cutedsl_multiple_padding_outputs_and_states: sentinel-filled padded output intervals and state slots exercised through the real wrapper, in eager and CUDA-graph replay, with both snapshot and ReplaySSM-ring state arms and both raw and fused-onorm outputs; real-request outputs must match a real-requests-only control run bitwise and all untouched padding state must stay unchanged. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> * test: cover one, two, and three padded requests in the KDA CuTeDSL multiple-padding test Extends test_cutedsl_multiple_padding_outputs_and_states to the three controlled conditions of that measurement: 7/6/5 real of 8 padded rows on a block-8 shape (npad 1/2/3), keeping the original 2-of-4 small-shape case. The npad-1 arm guards the already-correct single-padding geometry against regression; the npad>=2 arms fail on the unpatched kernel, which is what the measurement observed (first padded interval always zero, later padded intervals carrying torch.empty bytes). Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../kernels/ops/attention/kda_decode_mtp.py | 5 +- .../test_kda_mtp_cutedsl_replayssm_ring.py | 132 +++++++++++++++++- 2 files changed, 135 insertions(+), 2 deletions(-) diff --git a/python/sglang/kernels/ops/attention/kda_decode_mtp.py b/python/sglang/kernels/ops/attention/kda_decode_mtp.py index b3485b54bf98..ba617551d47b 100644 --- a/python/sglang/kernels/ops/attention/kda_decode_mtp.py +++ b/python/sglang/kernels/ops/attention/kda_decode_mtp.py @@ -272,7 +272,10 @@ def kda_decode_mtp_kernel( # CUDA-graph padding rows use slot == -1. if slot < 0: cute.arch.griddepcontrol_wait() - pad_bos = cu_seqlens[i_n] + # This wrapper requires a fixed dense verify width. Graph metadata + # collapses all padded sequence boundaries to the real-token endpoint, + # but each padded request still owns its own dense output interval. + pad_bos = i_n * T_LOOP for i_t in cutlass.range_constexpr(T_LOOP): if tidx < HEAD_DIM: o[0, pad_bos + i_t, i_hv, tidx] = cutlass.BFloat16(0.0) diff --git a/test/registered/kernels/ops/attention/test_kda_mtp_cutedsl_replayssm_ring.py b/test/registered/kernels/ops/attention/test_kda_mtp_cutedsl_replayssm_ring.py index 86bb5afb3873..39324269d55d 100644 --- a/test/registered/kernels/ops/attention/test_kda_mtp_cutedsl_replayssm_ring.py +++ b/test/registered/kernels/ops/attention/test_kda_mtp_cutedsl_replayssm_ring.py @@ -15,6 +15,8 @@ in fp32 registers. """ +from unittest.mock import patch + import pytest import torch @@ -44,7 +46,7 @@ W = 4 # KERNEL_WIDTH -def _run(arm, *, N, H, num_spec, seed, onorm=False, pad_last=False): +def _run(arm, *, N, H, num_spec, seed, onorm=False, pad_last=False, kwargs_only=False): torch.manual_seed(seed) T = N * (1 + num_spec) num_slots = N + 2 @@ -113,12 +115,23 @@ def _run(arm, *, N, H, num_spec, seed, onorm=False, pad_last=False): ) if arm == "baseline": inter = torch.zeros(N, 1 + num_spec, H, K, K, device=DEV, dtype=torch.float32) + if kwargs_only: + return dict(intermediate_ssm=inter, **kwargs) out = fused_kda_decode_mtp_dspark(intermediate_ssm=inter, **kwargs) return out, dict(inter=inter, slots=slots, scratch=scratch, **norm) rawv = torch.zeros(num_slots, H, L, K, device=DEV, dtype=torch.bfloat16) rawk = torch.zeros_like(rawv) gring = torch.zeros(num_slots, H, L, K, device=DEV, dtype=torch.float32) betar = torch.zeros(num_slots, H, L, device=DEV, dtype=torch.float32) + if kwargs_only: + return dict( + intermediate_ssm=None, + replayssm_rawv=rawv, + replayssm_rawk=rawk, + replayssm_g=gring, + replayssm_beta=betar, + **kwargs, + ) out = fused_kda_decode_mtp_dspark( intermediate_ssm=None, replayssm_rawv=rawv, @@ -188,6 +201,123 @@ def test_cutedsl_cuda_graph_padding_slot_is_safe(N): assert torch.count_nonzero(ring[name][N]).item() == 0 +@pytest.mark.parametrize( + "real_requests,padded_requests,width,heads", + [ + (7, 8, 8, 12), + (6, 8, 8, 12), + (5, 8, 8, 12), + (2, 4, 4, 2), + ], + ids=[ + "k3-seven-of-eight-block8", + "k3-six-of-eight-block8", + "k3-five-of-eight-block8", + "two-of-four-block4", + ], +) +@pytest.mark.parametrize("arm", ["baseline", "ring"], ids=["snapshots", "ring"]) +@pytest.mark.parametrize("onorm", [False, True], ids=["raw", "onorm"]) +@pytest.mark.parametrize("use_graph", [False, True], ids=["eager", "graph"]) +def test_cutedsl_multiple_padding_outputs_and_states( + real_requests, padded_requests, width, heads, arm, onorm, use_graph +): + kwargs = _run( + arm, + N=padded_requests, + H=heads, + num_spec=width - 1, + seed=41, + onorm=onorm, + kwargs_only=True, + ) + real_tokens = real_requests * width + kwargs["ssm_state_indices"][real_requests:] = -1 + # Production graph padding repeats the last real cumulative sequence length. + # Input/output storage still has a dense fixed-width block for every request. + kwargs["cu_seqlens"][real_requests + 1 :] = real_tokens + persistent_names = ("cs_q", "cs_k", "cs_v", "recurrent_state") + scratch_names = ( + "intermediate_conv_q", + "intermediate_conv_k", + "intermediate_conv_v", + "intermediate_ssm", + ) + ring_names = ("replayssm_rawv", "replayssm_rawk", "replayssm_g", "replayssm_beta") + state_names = [ + name + for name in persistent_names + scratch_names + ring_names + if isinstance(kwargs.get(name), torch.Tensor) + ] + # Nonzero sentinels also detect erroneous writes to unused state slots. + for name in scratch_names + ring_names: + if isinstance(kwargs.get(name), torch.Tensor): + kwargs[name].fill_(3) + before = {name: kwargs[name].clone() for name in state_names} + control = { + name: value.clone() if isinstance(value, torch.Tensor) else value + for name, value in kwargs.items() + } + for name in ("x_q", "x_k", "x_v", "g", "beta", "onorm_gate"): + if name in control: + control[name] = control[name][:, :real_tokens].contiguous() + for name in ("ssm_state_indices", "intermediate_state_indices"): + control[name] = control[name][:real_requests].contiguous() + control["cu_seqlens"] = control["cu_seqlens"][: real_requests + 1].contiguous() + expected = fused_kda_decode_mtp_dspark(**control) + + original_empty_like = torch.empty_like + + def sentinel_output(tensor, *args, **allocation_kwargs): + result = original_empty_like(tensor, *args, **allocation_kwargs) + if tensor is kwargs["x_v"]: + # Only intercept the wrapper's output allocation. The finite graph + # sentinel is refilled by a captured operation on every replay. + result.fill_(-77 if use_graph else float("nan")) + return result + + with patch.object(torch, "empty_like", sentinel_output): + if use_graph: + fused_kda_decode_mtp_dspark(**kwargs) + torch.cuda.synchronize() + for name in state_names: + kwargs[name].copy_(before[name]) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = fused_kda_decode_mtp_dspark(**kwargs) + graph.replay() + graph.replay() + else: + actual = fused_kda_decode_mtp_dspark(**kwargs) + torch.cuda.synchronize() + + torch.testing.assert_close(actual[:, :real_tokens], expected, rtol=0, atol=0) + padding = actual[:, real_tokens:] + assert torch.isfinite(padding).all().item() + assert torch.count_nonzero(padding).item() == 0 + # Full-buffer comparison includes untouched padding scratch/ring slots. + for name in state_names: + torch.testing.assert_close(kwargs[name], control[name], rtol=0, atol=0) + for name in persistent_names: + torch.testing.assert_close(kwargs[name], before[name], rtol=0, atol=0) + for name in scratch_names: + if name in before: + torch.testing.assert_close( + kwargs[name][real_requests:], + before[name][real_requests:], + rtol=0, + atol=0, + ) + for name in ring_names: + if name in before: + torch.testing.assert_close( + kwargs[name][real_requests + 1 :], + before[name][real_requests + 1 :], + rtol=0, + atol=0, + ) + + if __name__ == "__main__": import sys