get_device_context() allocates a new tensor on every call

[aamarnat]

Bug

In iris/iris.py, get_device_context() calls torch.tensor(context_data, dtype=torch.int64, device=self.device) every time it is invoked. The context data (rank, world_size, heap_bases) never changes after initialization.

This causes two problems:

1. Performance: Unnecessary GPU memory allocation on every kernel call.
2. CUDAGraph/HIPGraph incompatibility: torch.tensor() allocates GPU memory, which is forbidden during graph capture. Any kernel that calls get_device_context() in its launch path cannot be captured in a CUDAGraph.

Impact

Blocks CUDAGraph integration for any iris-backed fused kernel (e.g., matmul_all_reduce used via vLLM's torch.compile pipeline). Minor performance overhead otherwise.

Fix

Cache the tensor on first call:

def get_device_context(self):
    if self._cached_device_context is not None:
        return self._cached_device_context
    # ... existing construction logic ...
    self._cached_device_context = context_tensor
    return context_tensor

Add self._cached_device_context = None in __init__.

Component

iris/iris.py

[mawad-amd]

@copilot Fix this by precomputing the context tensor so get_device_context() is a simple return self._device_context with no allocation or branching.

Implementation:

1. Add a _build_device_context() private method to iris class in iris/iris.py that contains the current body of get_device_context() (lines 933-968). It should build the tensor and store it as self._device_context.

2. Call _build_device_context() at the end of __init__, after self.tracing = Tracing(self) is initialized.

3. Replace get_device_context() body with just return self._device_context.

4. In iris/tracing/core.py, at the end of Tracing.enable(), call self.iris._build_device_context() to rebuild the tensor with the tracing fields included.

Do not handle stale heap_bases refresh — that is tracked separately in #467.