Investigate `torch.Tensor.__setitem__` as unsupported in `thunderfx`

[mattteochen]

Setting this operator in

lightning-thunder/thunder/dynamo/utils.py

Line 48 in f6f0496

UNSUPPORTED_THUNDER_FUNCTION = ()

gives 10x decode throughput improvement on openai/gpt-oss-20b and Qwen/Qwen3-32B inside SGLang bench_one_batch.py.

Repro (this only works on the NV internal thunder-sglang-integration codebase):

SGLANG_USE_THUNDER_GRAPH_RUNNER=1 python3 -m sglang.bench_one_batch   --model-path openai/gpt-oss-20b --trust-remote-code   --model-impl transformers --dtype bfloat16   --json-model-override-args '{"quantization_config": null}' --cuda-graph-bs 1 --tp-size 4 --tp-strategy dtensor --load-format-dummy

[kiya00]

It also fits my observation when running the benchmark on each graph:

  Graph Name: graph7_thunder_1
    Operators used in the graph:
      <built-in function setitem>

the benchmark script can be a minimal repro:

from math import inf
from math import nan
NoneType = type(None)
import torch
from torch import device
import torch.fx._pytree as fx_pytree
import torch.utils._pytree as pytree
import torch
from thunder.dynamo.benchmark_utils import TorchInductorSpecification
from thunder.dynamo.benchmark_utils import KernelTime

from thunder.dev_utils import nvtx_profile_transform


class DynamoModule(torch.nn.Module):
    def forward(self, l_forward_batch_token_to_kv_pool_k_buffer_0_: "bf16[2615232, 8, 64][512, 64, 1]cuda:0", l_forward_batch_out_cache_loc: "i64[1][1]cuda:0", l_k_: "bf16[1, 8, 64][512, 64, 1]cuda:0"):
         # File: /sgl-workspace/sglang/python/sglang/srt/mem_cache/memory_pool.py:803 in set_kv_buffer, code: self.k_buffer[layer_id - self.start_layer][loc] = cache_k
        l_forward_batch_token_to_kv_pool_k_buffer_0_[l_forward_batch_out_cache_loc] = l_k_;  setitem = l_forward_batch_token_to_kv_pool_k_buffer_0_;  l_forward_batch_token_to_kv_pool_k_buffer_0_ = l_forward_batch_out_cache_loc = l_k_ = setitem = None
        return ()

def test_graph7_thunder_1():
    inputs = [
        torch.testing.make_tensor((2615232, 8, 64), dtype=torch.bfloat16,  device='cuda:0', requires_grad=False, low=-0.0013427734375, high=0.00133514404296875,),
        torch.full((1,), 0, dtype=torch.int64, device='cuda:0', requires_grad=False, layout=torch.strided),
        torch.testing.make_tensor((1, 8, 64), dtype=torch.bfloat16,  device='cuda:0', requires_grad=False, low=None, high=None,),
    ]

    model = DynamoModule()

    compiled_model = TorchInductorSpecification.torch_inductor(model, inputs, skip_symbolic_trace=True)

    #import thunder
    #compiled_model = thunder.jit(model,fusion_type='dataflow',thunderfx_disable_split_autograd=True,disable_torch_autograd=True,)

    # forward
    fwd_measurement = KernelTime(threshold=0.1, min_run_time=0.01, max_run_time=10.0).time("compiled_model(*inputs)", globals={"compiled_model":compiled_model, "inputs": inputs})
    print("fwd_measurement=", fwd_measurement)
    if hasattr(fwd_measurement, "max_allocated_memory"):
        from thunder.dynamo.benchmark_utils import get_pretty_memory_str
        print(f"fwd_measurement.max_allocated_memory={get_pretty_memory_str(fwd_measurement.max_allocated_memory)}")


if __name__ == "__main__":
    test_graph7_thunder_1()


"""
Environment information get from `torch.utils.collect_env.get_pretty_env_info()`:
CUDA devices:
  0: NVIDIA GB200
  1: NVIDIA GB200
  2: NVIDIA GB200
  3: NVIDIA GB200
CUDA version: 13.0
mypy_extensions==1.1.0
numpy==2.3.4
nvidia-cublas==13.0.0.19
nvidia-cuda-cupti==13.0.48
nvidia-cuda-nvrtc==13.0.48
nvidia-cuda-runtime==13.0.48
nvidia-cudnn-cu13==9.13.0.50
nvidia-cudnn-frontend==1.16.0
nvidia-cufft==12.0.0.15
nvidia-curand==10.4.0.35
nvidia-cusolver==12.0.3.29
nvidia-cusparse==12.6.2.49
nvidia-cusparselt-cu13==0.8.0
nvidia-nccl-cu13==2.27.7
nvidia-nvjitlink==13.0.39
nvidia-nvtx==13.0.39
optree==0.18.0
torch==2.9.0+cu130
torch_memory_saver==0.0.9
torchao==0.9.0
torchaudio==2.9.0
torchvision==0.24.0
triton==3.5.0

Versions of Thunder related libraries:
lightning-thunder==0.2.7.dev0
nvfuser==0.2.34+git569c115


"""

kernel time results:

thunder:
  Median: 1.70 ms
  IQR:    0.00 ms (1.70 to 1.70)
  6 measurements, 1 runs per measurement, 1 thread

inductor:
  Median: 1.56 us
  IQR:    0.01 us (1.56 to 1.56)
  65 measurements, 100 runs per measurement, 1 thread

[mattteochen]

The issue is how thunder decomposes a[pos] = value into:

b = a.clone()
b[pos] = value
a.copy_(b)

For small tensors, this doesn't affect overall performance, but for an inference workload with a large KV cache, this is very expensive due to GB of memory traffic to update a single index:

    def forward(self, l_forward_batch_token_to_kv_pool_k_buffer_0_: "bf16[2615232, 8, 64][512, 64, 1]cuda:0", l_forward_batch_out_cache_loc: "i64[1][1]cuda:0", l_k_: "bf16[1, 8, 64][512, 64, 1]cuda:0"):
         # File: /sgl-workspace/sglang/python/sglang/srt/mem_cache/memory_pool.py:803 in set_kv_buffer, code: self.k_buffer[layer_id - self.start_layer][loc] = cache_k
        l_forward_batch_token_to_kv_pool_k_buffer_0_[l_forward_batch_out_cache_loc] = l_k_;  setitem = l_forward_batch_token_to_kv_pool_k_buffer_0_;  l_forward_batch_token_to_kv_pool_k_buffer_0_ = l_forward_batch_out_cache_loc = l_k_ = setitem = None
        return ()

becomes:

# Constructed by Unwrap the actual return value
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(l_forward_batch_token_to_kv_pool_k_buffer_0_, l_forward_batch_out_cache_loc, l_k_):
  # l_forward_batch_token_to_kv_pool_k_buffer_0_: "cuda:0 bf16[2615232, 8, 64]"
  # l_forward_batch_out_cache_loc: "cuda:0 i64[1]"
  # l_k_: "cuda:0 bf16[1, 8, 64]"
  (t5,) = update_aliases((l_forward_batch_token_to_kv_pool_k_buffer_0_,))
  del l_forward_batch_token_to_kv_pool_k_buffer_0_

  # /kaiximatteoc/dev/sglang-thunder/python/repro.py:18:                l_forward_batch_token_to_kv_pool_k_buffer_0_[l_forward_batch_out_cache_loc] = l_k_;  setitem = l_forward_batch_token_to_kv_pool_k_buffer_0_;  l_forward_batch_token_to_kv_pool_k_buffer_0_ = l_forward_batch_out_cache_loc = l_k_ = setitem = None
  t6 = copy_with_setitem_impl(t5, l_forward_batch_out_cache_loc, l_k_)  # t6: "cuda:0 bf16[2615232, 8, 64]"
  [t1] = nvFusion0(t6, t5)
    # t1 = prims.copy_(t6, t5, grad_enabled=True)  # t1: "cuda:0 bf16[2615232, 8, 64]"
  del t6, t5
  return ()

Without considering host overhead to launch multiple kernels (memcpy and nvFuser pointwise), inductor's device speedup is > 1000x.

[mattteochen]

Thunder will face this same issue when compiling SGLang MoE implementation as in-place scatter and setitem are being used: https://github.com/sgl-project/sglang/blob/13680e5542848a6d33b436e4821e868476d060cd/python/sglang/srt/layers/moe/fused_moe_native.py#L49