Study Memo | 2026–07–8,9,10,11 | P8,12 Remove Softmax Dead Code, Generalize the Implementation, Strengthen the Tests/Fix the Guard Lifetime Issue in Asynchronous Execution

Diary

• I had worked hard on the concepts before, but I did not like reading code, so I only studied it superficially. After I got sick, took a break, and came back, I realized that trying to study without fully understanding things actually took even more time because I could not concentrate. This time I decided to study the code properly. After doing that, I could roughly understand what I was reading, and reading code did not take nearly as much time as I had expected. That made me really happy. Even now, though, when I understand something through written explanations, it feels like I understand every detail. When I read code, I spend so much energy on following it that I feel I understand the underlying principles a little less deeply. If I could understand everything completely, I think I would be much better at recognizing what I did wrong. I need to keep working on it.
• Telling myself that I would look at something later really meant I would never look at it. I had copied and pasted things before with the intention of reviewing them later, but I never went back to any of them. They just kept piling up, and eventually I could not even remember what was what. When organizing everything feels overwhelming, I should at least read through the important parts instead of saving them for later. From now on, I should stop thinking that I will come back to them later.
• I waste too much time. When I study by typing at my computer, the posture becomes uncomfortable. Whenever I take a short break, I end up watching TV or using my phone, and before I know it, a lot of time is gone. When I cannot stay at my desk, I think I should at least keep studying in my head while lying down or walking instead of wasting that time. The problem was that there was no book for what I am studying now, so I did not know what to review. Because of that, I asked AI to generate a large collection of study materials that I could use during short breaks. I want to keep studying even while resting so I can improve my skills as quickly as possible.

<July 8, 2026>

• boundary: The point where the compiler’s behavior changes during code execution.
• override: 1. Redefining inherited behavior. 2. Explicitly overriding or disabling a compiler’s default constraints or rules by applying an exception.
• guard: A conditional check that verifies whether it is safe to execute a particular piece of code.
• ComputeGloballyUnfusible is fine, but the ShouldFuse guard and its consumer-matching helper functions are unnecessary.
• inline: 1. Writing if (CheckShouldFuse(...)) is not inlining because the implementation is hidden inside another function. Writing the full condition directly, such as if (size > 0 && type == FLOAT && ...), is true inlining. 2. A comment left directly on a line of code in GitHub.
• Remove the guard and verify that everything still works by running the JAX reproduction or unit tests.
• Unit test: Compiles and tests only a small portion of the C++ code.
• Bazel test: Compiles the entire XLA C++ codebase and verifies that everything works together without errors.
• End-to-end test: Runs the full workflow by compiling the JAX Python code in a real user environment.
• Rebuilding jaxlib again.
• A positive test checks conditions where the bug should occur. A negative test checks conditions where the bug should not occur.
• In the negative test, verify that the override condition is not triggered.
• A few lines exceed 80 columns: Some lines are longer than 80 characters and need to be wrapped.
• clang-format: A tool for automatically formatting C++ code.
• Summary of tasks. First, remove the ShouldFuse guard and the consumer-matching helper functions. Second, run unit tests and test with jaxlib, then add a negative test. Third, run clang-format to clean up the code.

inline comment1

// Fusing one side of a reduce-max / exp-shift coupled producer into a
// separate kernel can recompute the shift value with different numerics.
if (HasCoupledReductionShiftExpSplitRisk(producer, consumer) &&

• This code will be removed.
• producer: The upstream operation that contains the reduce-max.
• consumer: The downstream exp operation that consumes its output.
• These two operation nodes are checked to determine whether splitting them could introduce numerical differences that result in NaN values.
• Delete all helper functions used by ShouldFuse, including HasCoupledReductionShiftExpSplitRisk, IsReduceSideConsumer, IsShiftExpSideConsumer, FusionLooksLikeShiftExpPath, and InstructionOrFusionContains.
• Keep only the IsCoupledReductionShiftExpProducer matcher and the override.
• The multiply operation is already blocked by do_not_duplicate.
• In an expression with two operands and one operator, the multiplication symbol is the operator, and the values on each side are the operands.
• Constants and enum values use the prefix k. In XLA, HLO opcodes are defined in the HloOpcode enum, where matrix multiplication is represented by kDot. Multiply performs element-wise multiplication between tensors of the same shape, while kDot represents standard matrix multiplication.
• kDot is already prevented from being fused because it is a large operation and is filtered out by CanBeLoopFused(*producer).
• Fusion interior refers to the operations inside a fused kernel. A FusionInstruction represents a kernel composed of multiple fused operations. Fusion exterior refers to the fused kernel itself when viewed from the outside, while fusion interior refers to the operations inside it.
• Due to the compiler’s structure, matchers cannot match operations inside the fusion interior.
• Fused parameters are the inputs that allow values from outside the fused kernel to be passed in. They use kParameter to receive values produced by other kernels.
• Without fusion, matching is possible. After fusion, however, the operands of the reduce-max become fused parameter nodes instead of operation nodes, so the matcher can no longer match them.
• A pointer stores a memory address. Pointer comparison checks whether two pointers refer to the exact same memory location rather than whether their values are equal.
• A predicate is a function that returns either true or false.
• Broadcast expands a smaller value or tensor so it can be used with a larger tensor shape.
• Cheap operations are often duplicated instead of being stored in memory. Because of this duplication, the operations became split, which eventually caused the NaN bug.
• Remove all unnecessary code, then run the unit tests, rebuild jaxlib, reproduce the issue, and add a negative test.

inline comment2

// A multiply that feeds both row-max reduction and exp-shift subtraction.
bool IsCoupledReductionShiftExpProducer(const HloInstruction* instr) {
if (instr->opcode() != HloOpcode::kMultiply) {

📝 HloInstruction is the class that represents a single operation node.

🌱 * indicates a pointer, which stores a memory address.

🌼 const means the object will only be read and will not be modified.

🐣 In instr->, the -> operator is used to access a member function through a pointer. opcode() returns the operation type, such as addition or subtraction. instr->opcode() asks what kind of operation the current instruction represents.

🍀 :: is the scope resolution operator and means "belongs to." HloOpcode::kMultiply refers to the kMultiply value defined in the HloOpcode enum, which represents element-wise multiplication.

🌸 Not only kMultiply, but other inexpensive operations such as addition and subtraction can also be treated as EffectivelyAtMostUnary, allowing them to be split. That can also lead to the same NaN issue.

🌿 In the expression where an input is multiplied by a weight matrix and then a bias is added, the bias is represented by b.

✨ add(dot, bias) represents adding the bias after a matrix multiplication.

🍃 Generalize the pattern as much as possible without introducing numerical correctness issues.

inline comment3

EXPECT_FALSE(HasSplitCoupledReductionShiftExpFusion(*module));
}

TEST_F(InstructionFusionTest, StillFusesUncoupledReduceAndSubtract) {

🌼 EXPECT_FALSE(...) passes only if the expression inside evaluates to false. HasSplitCoupledReductionShiftExpFusion checks whether the softmax-related operations have been split.

🌱 module is the complete computation graph to be compiled. When user-written JAX code is passed to the compiler, it is represented as a single large HloModule.

📝 TEST_F is a macro provided by GoogleTest. The F stands for fixture, meaning the test runs with a predefined test environment. InstructionFusionTest provides that fixture, and StillFusesUncoupledReduceAndSubtract is the name of this unit test.

🌸 “It doesn’t exercise the matcher’s discrimination” means the test does not actually verify whether the matcher can distinguish the bug pattern from valid patterns.

🍀 “Trivially rejects everything” means everything is rejected immediately by an early condition.

🐣 The test module that I created does not contain any multiply operations. Since the matcher begins by checking whether the opcode is HloOpcode::kMultiply, every instruction is rejected immediately, so the later softmax structure is never examined.

🌿 “And the assertion is only changed == true, which any fusion in the module satisfies” means the assertion succeeds whenever any fusion occurs, regardless of whether the intended pattern was matched.

✨ Add at least one kMultiply operation to the HLO graph. That allows the matcher to reach the softmax structure, where it can verify that non-bug patterns are still fused correctly.

🌼 The test should also include decoy patterns that are almost identical to the bug pattern but differ in one essential way, verifying that the matcher correctly classifies them as valid.

🌱 A multiply feeding an add-reduce plus subtract means that the multiply branches into an addition-based reduction on one side and a subtraction on the other. Since this is not a max-reduction, it is not the bug pattern and should still be fused.

🍃 A max-reduce whose subtract does not feed an exp means the graph matches the bug pattern up to the subtraction, but the subtraction is not followed by an exponential operation. This should also remain fusible.

🌸 The final entry computation produced by the compiler should not contain any standalone multiply operations.

🐣 Fusion is the compiler’s default behavior because it is generally more efficient.

🌱 Remove the ShouldFuse guard and the consumer-matching helper functions.

🌼 Generalize the implementation as much as possible without introducing numerical correctness issues.

🍀 Run clang-format to wrap overly long lines and clean up the formatting.

🌿 Update the test so that it begins with a kMultiply operation, then verify that both "a multiply feeding an add-reduce plus subtract" and "a max-reduce whose subtract doesn't feed exp" are still fused correctly.

✨ Run the unit tests and jaxlib tests, and add negative tests as described above.

📝 HloInstruction is the class that represents a single operation node.

🌱 ; marks the end of a statement in C++.

🌼 else if means "otherwise, if." It is evaluated only when the preceding if condition is false.

🐣 -> is used to call a member function through a pointer.

🍀 opcode() returns the opcode that identifies the operation type.

🌸 In for (...), the expression after the colon specifies the range to iterate over.

🌿 || is the logical OR operator. It evaluates to true if either side is true.

✨ A .h header file declares the functions and interfaces implemented in the corresponding .cc source file.

📝 If functions are defined only inside an anonymous namespace in a .cc file, they have internal linkage and are not declared in the .h file because they are not used outside that translation unit.

🌱 absl::c_any_of(...) returns true if any element satisfies the given condition.

🌼 [&] captures variables from the surrounding scope by reference.

🐣 LoopFusion represents a fused computation rather than a single operation.

🍀 absl::string_view module_string = R"(\nHloModule module\n...)" stores an HLO module as readable text in module_string.

🌸 ParseAndReturnVerifiedModule(...) parses the textual HLO module, verifies it, and creates an HloModule object.

🌿 I understand the idea of generalizing the implementation to element-wise operations, but I am still unsure whether matching only the softmax pattern is sufficient. Softmax itself begins with a multiply, but in real LLMs and Transformer models, additional operations are usually inserted before softmax. Common examples include add(dot, bias). Before the dot product, the query and key are multiplied. After that, a bias is added, then an attention mask is added to mask specific positions, and only then is softmax applied.

✨ When creating a test HLO module, I start with declarations such as f32[8,16] parameter(0) and then assign names like %arg = ... or %reduce = .... The % prefix is used to name values in HLO.

📝 Checking the HLO dump is useful because it shows the compiler’s intermediate representation, making it possible to see the actual computation graph produced during compilation.

🌱 absl::c_linear_search performs a linear search through a container from beginning to end. The c stands for container.

🌼 The standard C++ approach requires explicitly passing iterators, such as std::vector<int> numbers = {1, 2, 3, 4, 5}; bool found = std::find(numbers.begin(), numbers.end(), 3) != numbers.end();. By including #include "absl/algorithm/container.h", the same search can be written as std::vector<int> numbers = {1, 2, 3, 4, 5}; bool found = absl::c_linear_search(numbers, 3);, which is more concise.

🐣 %max (p0: f32[], p1: f32[]) -> f32[] {
%p0 = f32[] parameter(0)
%p1 = f32[] parameter(1)
ROOT %m = f32[] maximum(%p0, %p1)
}

This is the textual HLO representation that the XLA compiler uses internally. % introduces the name of an HLO value or computation. %max is the computation name, and it takes two input parameters inside the parentheses. Both p0 and p1 are scalar values of type f32. The arrow indicates that the computation returns a scalar f32. %p0 = stores the first input parameter in the local HLO value named p0, and %p1 does the same for the second input. ROOT marks the final result of the computation. %m is the value that holds the final result, and maximum(%p0, %p1) returns the larger of the two input values.

More Generalization

🌱 Inexpensive operations such as addition and subtraction, in addition to kMultiply, can also be classified as EffectivelyAtMostUnary, making them eligible for duplication during fusion. The same kind of NaN issue can occur after they are duplicated and split.

🌼 In the expression where an input is multiplied by a weight matrix and then a bias is added, b represents the bias.

🐣 add(dot, bias) represents adding a bias to the result of a matrix multiplication.

🍀 Generalize the implementation as much as possible without introducing numerical correctness issues.

🌸 A C++ unit test does not use a CPU, GPU, JAX, or real execution. It simply creates a small HLO graph in text form, runs it through the modified compiler code, and checks how the graph is fused. It does not verify numerical correctness or whether NaN values occur.

🌿 A JAX reproduction test is performed after rebuilding jaxlib with the modified compiler code. After rebuilding, the updated jaxlib is imported into Python, dummy data is passed through repro.py, and the program is executed on a GPU or TPU to verify that the NaN issue has actually been eliminated.

✨ Add the negative test cases discussed earlier to the unit tests and verify that they all behave as expected.

📝 opcode identifies the exact operation type, such as kMultiply or kAdd.

🌱 Topology describes how nodes are connected to one another in the computation graph.

🌼 Blast radius refers to the range of code that could be affected by a change.

🐣 If opcode() != HloOpcode::kMultiply is replaced with an element-wise check, it can match operations such as kAdd and kSubtract as well. Because the surrounding conditions still filter the matched cases, the risk of affecting unrelated code should remain low.

🍀 kFusion represents a fused operation that contains multiple operations.

🌸 post_order is the traversal order that visits the graph from the outputs back toward the inputs.

🌿 The compiler performs fusion in multiple iterations.

✨ The current issue is that inexpensive operations are considered safe to duplicate simply because they are cheap. However, when a kFusion contains only inexpensive operations, it can also be treated as an inexpensive operation internally. Because of that, using opcode != kFusion && IsElementwise() instead of checking only for kMultiply can catch more problematic cases while still excluding fusion instructions.

Test Cases

📝 There are currently three tests: FusionComputesShiftValueForReduce, the FusionRecomputesShiftValueForExpShift helper, and the positive and negative tests.

🌱 In the current tests, the risky starting point, shift_value, receives a raw external parameter %arg instead of the result of an operation. As a result, there is no multiply operation, so the matcher exits before reaching the later pattern. The test passes, but it does not actually verify anything useful.

🌼 changed becomes true whenever any fusion modifies the graph. What I actually want to verify is whether the reduce and subtract operations were fused together. Since unrelated fusion elsewhere also makes changed true, this is not a reliable test. Instead, the tests should directly inspect the resulting HLO graph.

Things to Fix

– Unit Tests

🌱 Keep AvoidSplitCoupledReductionShiftExpFusions as it is. It verifies that the coupled softmax pattern is not split.

🌼 Make the negative tests more precise. They should verify that slightly modifying the bug pattern causes the matcher to correctly recognize it as a negative case.

🐣 Instead of checking only changed == true, verify that no standalone multiply operation remains in the entry computation.

🍀 Since the matcher now uses IsElementwise, add another positive test in which the shift value is produced by add(dot, bias) instead of multiply(dot, scale).

– Rebuild jaxlib and Run repro.py

🌸 Reproduce the issue using the scaled softmax pattern on the CPU backend. The pattern should be softmax(scale * (q @ k.T)), which is the same shape used by the original reproduction. Run it with the rebuilt jaxlib.

🌿 Compare the results against a high-precision reference, such as float64 or NumPy, or against another backend where the split does not occur. Verify the outputs using np.allclose with a strict tolerance. Before the fix, the results should differ or produce NaN values. After the fix, they should match.

✨ To verify the generalized implementation, also run the softmax(q @ k.T + bias) version alongside the corresponding unit test so that the element-wise generalization is confirmed end to end.

📝 Optionally, generate an HLO dump using XLA_FLAGS=--xla_dump_to=.... Check that the coupled multiply or add is not split into two separate fusion kernels and confirm that both the reduce path and the exp-shift path still share the same shift value. This is also a good opportunity to make sure that unrelated fusion behavior has not changed unexpectedly.

🌱 Run both the unit tests and the jaxlib tests on a Google Cloud VM instance. The HLO dump should also be generated there.

🌼 Run clang-format.

<Thursday, July 9, 2026>

📝 to_apply() retrieves the computation that will be applied internally during an operation.

🌱 root_instruction() returns the root instruction of a computation.

🌼 & represents a memory address.

🐣 In softmax-related code, shift_value often refers to the scaled value.

🍀 A lambda function consists of captures inside [], parameters inside (), and the implementation inside {}.

🌸 [&] captures all surrounding variables by reference.

🌿 absl::Span<...> provides an efficient view over existing memory without copying it. The elements inside the span are the values being examined.

✨ HloReachabilityMap& stores reachability information so it is already known whether one instruction can reach another.

📝 #include includes a header file.

🌱 A function must be declared in the .h file before code in other translation units can call its implementation in the .cc file.

🌼 override indicates that a virtual function inherited from the base class is being overridden.

🐣 In TEST_F, the arguments inside the parentheses define the test fixture and the test name. The first argument specifies the test environment, and the second is the name of the individual test.

🍀 contracting_dims specifies which dimensions are contracted during matrix multiplication.

🌸 to_apply=%max means that the previously defined maximum computation is used for the reduction.

🌿 dimensions={1} means the reduction is performed along dimension 1 to compute the maximum values.

✨ ParseAndReturnVerifiedModule(module_string) parses the textual HLO module, verifies it, and creates an HloModule.

📝 TF_ASSERT_OK_AND_ASSIGN(auto module, ...) stops immediately if an error occurs. Otherwise, it stores the resulting module in module.

🌱 .Run(module.get()) runs the optimization pass on the module using the modified compiler.

🌼 bool changed records whether the optimization pass modified the original module.

🐣 negate changes the sign of a value.

🌿 Generate an HLO dump while running the unit tests, then rebuild jaxlib, and finally run the reproduction test.

Unit Tests + HLO Dump

📝 Connect to the VM through SSH.

🪴 Clone the XLA repository from my branch.

git clone https://github.com…

cd xla

git checkout avoid-split-softmax

git pull origin avoid-split-softmax

🌸 Set up the build environment.

sudo apt-get update

sudo apt-get install -y docker.io git

sudo usermod -aG docker $USER

🌱 Install Docker and Git so the project can be built safely in an isolated environment. Large open-source projects are typically built inside Docker containers on Linux systems to ensure a consistent build environment.

🍀 sudo runs a command with administrator privileges. apt-get update refreshes the package index.

🐣 docker.io installs the Docker Engine.

🌼 By default, Docker usually requires elevated permissions. This can cause permission issues when running build scripts or when Bazel communicates with Docker. usermod modifies the current user's account settings. -aG docker adds the current user to the Docker group, and $USER refers to the currently logged-in user.

cd ~/xla

docker run -itd — rm — name xla -w /xla -v $PWD:/xla
us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest
bash

docker exec xla ./configure.py — backend=CPU

🌿 docker run starts a new Docker container. In -itd, the d option runs the container in detached mode so it continues running in the background. --rm automatically removes the container after it stops. --name xla assigns the container the name xla.

Get Lia’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

✨ -w /xla sets /xla as the working directory inside the container.

📝 -v $PWD:/xla mounts the current directory from the VM into the container so both environments share the same files.

🌱 us-docker.pkg.dev/.../ml-build:latest is the Google-maintained build image that already contains compatible versions of Bazel, Python, and the C++ toolchain.

🌸 docker exec xla runs a command inside the running container.

Unit Tests

🪴 (1) PR-related tests: Verify that the softmax bug has been fixed.

bazel test -c opt — test_output=errors
//xla/service/cpu:cpu_instruction_fusion_test
— test_filter=”AvoidSplitCoupled:StillFusesMultiply”

🍀 -c specifies the compilation mode. opt builds an optimized version. --test_output=errors prints only the logs from failed tests.

🌼 (2) Fusion tests: Verify that the changes do not affect unrelated code and that the full fusion test suite still passes.

bazel test -c opt — test_output=errors
//xla/service/cpu:cpu_instruction_fusion_test

HLO Dump

🐣 Save the HLO used in the unit test to a file, then run run_hlo_module.

🌿 Create the HLO file.

mkdir -p /tmp/hlo

nano /tmp/hlo/softmax_scale.hlo

✨ When the Nano editor opens, paste the HLO from the AvoidSplitCoupledReductionShiftExpFusions unit test. Nano is a text editor commonly used in Linux environments.

📝 Run the HLO dump.

mkdir -p /tmp/xladump

bazel run -c opt //xla/hlo/tools:run_hlo_module —
— platform=CPU
/tmp/hlo/softmax_scale.hlo
— xla_dump_to=/tmp/xladump
— xla_dump_hlo_as_text
— xla_dump_hlo_pass_re=fusion

🌱 A regular expression, or regex, is used to match only the text patterns that I want to search for.

🌸 Check the results.

ls /tmp/xladump

grep -l “multiply|fusion” /tmp/xladump/*.txt | head

<Friday, July 10, 2026>

📝 xla/backends/gpu/collectives/gpu_cliques.cc: In graph theory, a clique is a subgraph in which every vertex is connected to every other vertex. Here, it represents a GPU communication group. The GPUs participating in collective communication are managed as a single clique.

🌱 absl::Status is an object that represents the result of an operation. It indicates whether the operation succeeded or failed and, if it failed, contains information about the reason.

🌼 absl::MutexLock is an RAII lock class that automatically acquires and releases a mutex. state.mu is the shared mutex that protects shared state.

🐣 A vector is a dynamically sized array. std::vector<coordination::TaskInfo> means that the vector stores coordination::TaskInfo objects.

🍀 auto lets the compiler deduce the variable's type automatically.

🌸 it == updated.end() checks whether the iterator reached the end of the container without finding a matching element.

🌿 %d is the format specifier for an integer.

✨ it->set_state updates the state of the current object.

📝 it->mutable_error_payload()->set_source_task_id(failed_task_id); obtains a mutable pointer to the error payload and records the task ID where the error originally occurred.

🌱 state.task_state_infos = std::move(updated); transfers the contents of updated into state.task_state_infos without copying them.

🌼 Execution threads are responsible for issuing work to the GPU and coordinating data transfers.

🐣 The Hang Watchdog monitors execution threads. If a thread remains inactive for too long, it detects the hang.

🍀 If an execution thread is blocked while waiting for the GPU, it cannot respond even if cancellation is requested. The Hang Watchdog signals it through a CancellationToken, allowing the blocked operation to unwind and exit safely.

🌸 xla/pjrt/distributed/coordination/coordination_client.h: When multiple GPUs work together, they need to share their execution state. This client defines the interface used to communicate with the central coordination service.

🌿 A function declared as virtual ... = 0 is a pure virtual function. It provides only the interface, while derived classes implement the actual behavior.

✨ void means that the function does not return a value.

📝 Asynchronous error reporting sends an error notification and continues executing without waiting for a response.

🌱 tsl stands for the TensorFlow System Libraries namespace.

🌼 A .proto file defines Protocol Buffers schemas. It specifies how structured data is represented so that different programming languages can exchange the same data consistently.

protobuf Syntax

🐣 message defines a structured data type.

🍀 Incarnations incarnations = 1; declares a field named incarnations of type Incarnations. The field number 1 uniquely identifies it in the serialized Protocol Buffers format.

🌸 rpc stands for Remote Procedure Call. It allows programs running on different machines to communicate as though they were calling local functions.

🌿 absl::MutexLock l(state_mu_); creates a lock object that acquires the mutex named state_mu_.

✨ leader_client_ is the client object used to communicate with the leader node.

📝 In LOG(ERROR), the << operator appends data to the log message.

🌱 ExecuteThunksImpl is the function where GPU work is launched. A watchdog guard is created there to detect execution hangs.

🌼 The guard starts a timer. If the execution thread remains inactive for ten minutes, it reports a timeout.

🐣 The execution thread submits work to the GPU asynchronously. It issues the command and continues without waiting for the GPU to finish. The GPU executes the work independently in the background.

🍀 When a C++ function returns, its local variables are destroyed. The guard created inside the function is therefore a local variable.

🌸 A class is a blueprint used to define objects.

🌿 An object is a concrete instance of a class that exists in memory.

✨ Lifetime refers to the period from when an object is created until it is destroyed.

📝 A function is like a stage, and local variables are temporary props used only during that performance. Once the function ends, those props are removed. Member variables belong to an object itself. As long as the object remains alive, its member variables continue to exist, even after individual member functions return. They are destroyed only when the object itself is destroyed.

🌱 A local variable is created inside a function or block and is destroyed when the enclosing braces are exited.

🌼 A member variable belongs to a class instance and remains alive until the object itself is destroyed.

🐣 guard = HangWatchdog::Global().Watch(watchdog_name, watchdog_timeout, std::move(on_timeout));

🍀 HangWatchdog::Global() returns the global Hang Watchdog instance.

🌸 .Watch(...) starts monitoring the current operation.

🌿 std::move(on_timeout) transfers ownership of the timeout callback to the watchdog so it can invoke the callback if a timeout occurs.

📝 When the code above runs, HangWatchdog creates a guard. However, ExecuteThunksImpl is asynchronous. It only submits work to the GPU and then returns immediately. As soon as the function returns, the local guard object is destroyed, which automatically unregisters it from HangWatchdog, so the GPU is no longer being monitored.

🌱 A struct is almost the same as a class. It is a compound data type that groups related variables and data members together. The main difference in C++ is that members of a struct are public by default, while members of a class are private by default.

🌼 A pointer (*) stores the memory address of an object rather than the object itself.

🐣 std::shared_ptr is a smart pointer that allows multiple owners to share the same object. The object remains alive until the last shared_ptr that owns it is destroyed.

<Wrong solution>

🍀 class GpuExecutableRunOptions { public: // ... ...

private: bool requires_exclusive_lock_on_gpu_ = false;

bool enable_mock_collectives_ = false;

// The HangWatchdog guard should be stored here as a member variable so it remains alive.

ExecutionTimeoutHandler execution_timeout_handler_; };

🌸 GpuExecutableRunOptions is created before GPU execution begins and remains alive until the asynchronous GPU work finishes. If the guard is stored as a member variable of this object, it will continue monitoring the GPU even after ExecuteThunksImpl has returned because the GpuExecutableRunOptions object itself is still alive.

Current Code Status

🌿 xla/service/gpu/gpu_executable.cc is where GPU execution begins. This is where the watchdog guard is created and the GPU work is launched.

✨ guard = HangWatchdog::Global().Watch( watchdog_name, watchdog_timeout, HangWatchdog::Abort(watchdog_name, watchdog_timeout, std::move(pre_abort)));

HangWatchdog::CancelCallback on_timeout;

if (gpu_run_options && gpu_run_options->execution_timeout_handler())

{ on_timeout = [watchdog_name, watchdog_timeout, gpu_run_options]() mutable { gpu_run_options->execution_timeout_handler()(watchdog_name, watchdog_timeout); }; }

📝 [watchdog_name, ...]() mutable { ... } is a lambda function. It creates a callback that will be executed if the timeout occurs and stores it in on_timeout.

🌱 In xla/pjrt/gpu/se_gpu_pjrt_client.cc, the timeout handler is executed when the GPU stops responding.

🌼 gpu_run_options->set_execution_timeout_handler( [process_index = options.node_id, distributed_client = options.distributed_client] (absl::string_view action, absl::Duration timeout) { absl::Status error = absl::DeadlineExceededError( absl::StrFormat("%s failed to finish in %v", action, timeout)); if (absl::Status s = gpu::AbortCollectivesOnTaskFailure(process_index, error); !s.ok()) { LOG(WARNING) << s; } if (distributed_client != nullptr) { absl::StatusOr<CoordinationServiceAgent*> agent = distributed_client->GetCoordinationServiceAgent(); if (agent.ok()) { (*agent)->ReportError(error); } } });

🐣 In xla/pjrt/distributed/coordination/coordination_service_agent.cc, calling ReportError from the previous step reaches this file, where an RPC is sent to the coordination service.

🍀 `absl::Status CoordinationServiceAgent::ReportError(const absl::Status& error) {
{
absl::MutexLock l(state_mu_);
if (state_ == xla::coordination::TaskState::ERROR) return absl::OkStatus();
}

ReportErrorToServiceRequest request;
request.set_error_code(error.raw_code());
request.set_error_message(std::string(error.message()));
absl::Notification n;
leader_client_->ReportErrorToServiceAsync(&request, &response, [&](const absl::Status& s) {
n.Notify();
});
n.WaitForNotification();
return absl::OkStatus();
}`

🌸 The overall flow is as follows. gpu_executable.cc creates the watchdog guard and submits work to the GPU. If the GPU remains unresponsive for ten minutes, the timeout handler in se_gpu_pjrt_client.cc is triggered, which aborts the GPU collectives and reports the timeout. The error is then forwarded through a gRPC call to the coordination service, where the coordinator records the failure and notifies the rest of the distributed system.

📝 std::shared_ptr<HangWatchdog::Guard> guard = nullptr;

on_timeout = [..., &guard]() mutable { ... };

guard = HangWatchdog::Global().Watch(..., std::move(on_timeout));

return MaybeSyncAndProfile(..., block_host_until_done ? main_stream : ...);

🌱 std::shared_ptr<...> is a smart pointer that keeps the object inside it alive until the last owner releases it.

🌼 In on_timeout = [..., gpu_run_options, &guard]() mutable { ... }, the capture list inside [...] specifies which variables from the surrounding scope are available inside the lambda.

🐣 PJRT stands for Portable JIT Runtime. It is XLA’s runtime interface that allows machine learning frameworks to communicate with different hardware backends.

🍀 VMM stands for Virtual Memory Management. On recent NVIDIA GPUs, it provides more flexible management of GPU virtual memory.

🌸 An allocator is responsible for managing memory allocation and deallocation.

🌿 When block_host_until_done == false, the host CPU thread does not wait for the GPU work to finish.

✨ Enqueue means placing a task into a queue so that the GPU can execute it later.

📝 MaybeSyncAndProfile determines whether the host should synchronize with the GPU and whether profiling should be performed after the work has been submitted.

🌱 In this case, synchronization is not performed because the stream is asynchronous. The function returns immediately without waiting for the GPU to finish.

🌼 HangWatchdog continues monitoring only while the shared_ptr<Guard> remains alive. Internally, the watchdog keeps only a weak reference to the guard. Once the last shared_ptr is destroyed, the weak reference expires, and the watchdog automatically stops monitoring that execution.

🐣 As a result, the timeout callback, on_timeout, is never invoked.

🌸 NCCL, the NVIDIA Collective Communications Library, provides high-performance communication between GPUs and is a core component of distributed training.

🌿 all_reduce_thunk.cc contains the implementation that launches collective all-reduce operations across GPUs.

✨ future.Await() waits for an asynchronously launched operation to complete. A future represents the result of work that is still in progress, and calling .Await() blocks until that work finishes. In this case, it waits for the collective operation to complete.

📝 Because of this, if one rank becomes stuck during a collective operation, ExecuteThunksImpl does not return. The guard therefore remains alive, allowing the first Hang Watchdog to continue monitoring the execution.

🌱 Stack variables are allocated inside a function’s stack frame. When the function returns, the stack frame is destroyed, and all local variables, including guard, are destroyed as well.

🌼 After the initial timeout triggers a post-abort action, a second watchdog with a shorter timeout is armed to detect cases where the process remains hung even after the first abort. The re-arming logic also depends on the same stack-allocated guard.

🐣 Once collective cancellation completes and ExecuteThunksImpl returns, both the first and second watchdog guards are destroyed.

🍀 If one GPU fails, collective cancellation allows the remaining execution to unwind and terminate cleanly.

🌸 A GPU kernel is a single computation executed on the GPU.

<The reason why it’s wrong>

🌿 The client is the application that launches the computation. GpuExecutableRunOptions is a configuration object that exists for the lifetime of the program, while an execution represents a single GPU execution. Since a guard should exist only for the lifetime of one execution, storing it inside GpuExecutableRunOptions is not appropriate.

✨ GpuExecutableRunOptions is created once and shared throughout the lifetime of the application. If two GPU executions overlap, storing the watchdog guard there could cause the guard for one execution to be overwritten by another. When the first execution finishes and its timer is stopped, it could unintentionally affect the watchdog associated with the second execution.

📝 Layer coupling describes how strongly different software layers depend on one another. Lower coupling generally makes the system easier to maintain because changes in one layer are less likely to affect another.

🌱 GpuExecutableRunOptions belongs to the configuration layer, while HangWatchdog belongs to the execution layer. Mixing these responsibilities increases coupling and makes the design less clean.

<Finding the right solution>

🌼 Synchronous or blocking execution means the CPU waits until the requested work has finished before continuing.

🐣 The caller is the function that invokes another function. In this case, the caller is the higher-level system that invokes ExecuteThunksImpl. Instead of storing the guard as a stack variable or inside the configuration object, it would be better for ExecuteThunksImpl to own it through a std::shared_ptr.

🍀 Above gpu_executable.cc, asynchronous execution is orchestrated by higher-level functions such as RunAsync and ExecuteAsyncOnStreamImpl, which manage the overall execution flow.

🌸 A stream completion event signals that the GPU has finished all queued work on the stream. A buffer-ready future indicates that the output buffer is ready to be consumed. Keeping the guard alive until the stream completion event occurs would ensure that the watchdog monitors the entire execution.

🌿 Heap memory has a lifetime independent of the current function. Unlike stack memory, objects allocated on the heap remain alive until they are explicitly destroyed or until their owning smart pointers release them.

✨ Value capture copies a variable into the lambda object itself. The previous approach used &guard, which captured only the address of the local variable. If the stack frame was destroyed, that address became invalid. Capturing by value avoids that problem.

📝 A holder wraps the guard in an additional level of ownership, such as std::shared_ptr<std::shared_ptr<Guard>>. The outer shared_ptr lives on the heap and owns the inner shared_ptr, which in turn owns the guard.

🌱 The previous implementation captured &guard so that the lambda could later assign the result of Watch(...). Once the stack frame disappeared, that pointer became a dangling pointer. Replacing the reference capture with value capture avoids this issue.

🌼 By allocating the holder on the heap, the copied shared_ptr remains alive even after the function returns. The lambda captures the holder by value instead of capturing a stack variable by reference. Even after the stack variables are destroyed, the holder remains valid, avoids dangling references, and still allows the post-abort callback to update the guard after the timeout.

🐣 This approach makes the implementation somewhat more complicated. The real goal is to tie the guard’s lifetime directly to the lifetime of a single execution. ExecutionScopedState already represents exactly that concept. Since it exists for the duration of one GPU execution, storing the guard there naturally keeps it alive until the execution scope ends.

Current Limitations of the Tests

📝 The current tests do not verify whether the actual timeout path is executed. They only check whether the timeout handler is installed when abort_collectives_on_failure = true.

🌱 Instead of waiting for a real ten-minute timeout, the tests invoke the handler directly and verify that ReportError is sent to the coordination service.

🌼 They also verify that the code behaves correctly when the distributed_client is nullptr.

What the Current Tests Do Not Verify

🐣 They do not verify that ExecuteThunksImpl still behaves correctly when the xla_gpu_execution_terminate_timeout option is enabled, or that the watchdog continues monitoring after an asynchronous return without manually invoking the handler.

🍀 They do not verify that when one rank stalls, the remaining ranks are released from the NCCL collective and the distributed training job terminates correctly.

🌸 They do not verify that if the system remains hung after the first timeout, the post-abort fallback triggers a process abort one minute later.

Test Plan

🌿 Configure a short timeout of about two seconds. Keep the GPU work running by using a mock collective or a sleep thunk. Execute the block_host_until_done = false path so that the execution is fully asynchronous. After the work has been enqueued and ExecuteThunksImpl has already returned, verify that the timeout callback is still invoked. This can be implemented as a Bazel test.

✨ Create a two-node environment with abort_collectives_on_failure = true. Trigger a failure on Node 0 so that the timeout handler runs. Verify that Node 1 receives a DeadlineExceededError through PollForError after the coordination service broadcasts the error. Then verify that Node 1 invokes AbortCollectivesOnTaskFailure and that the mock NCCL implementation confirms the collective operation is unblocked.

📝 Test the post-abort fallback by intentionally making collective cancellation fail after the timeout handler has been invoked. Verify that a process abort occurs one minute later.

🌱 Slurm and Kubernetes distribute work across multiple machines and manage distributed jobs. An exit code indicates how a program terminated. An exit code of 0 means successful completion, while any nonzero value indicates an error or abnormal termination. When such a failure occurs, Slurm or Kubernetes can automatically restart the job. Nightly builds are automated builds that run every night.

xla/service/gpu/gpu_executable.cc

[Upper Layer: High-Level Execution]

↳ RunAsync() or ExecuteAsyncOnStreamImpl()

— — — — — — — — —(★ Boundary where the scope of the change expands)

[Lower Layer: Execution]

↳ ExecuteThunksImpl() ← This is the main function I am currently working on.

🪴 The first watchdog timer is created at the beginning of ExecuteThunksImpl. A HangWatchdog::Scope or HangWatchdog::Guard object is created as a local stack variable.

🌸 Timeout handler:

void GpuExecutable::on_timeout() {

// 1. Report the failure to the coordination service.

coordination_agent_->ReportError(errors::DeadlineExceeded("GPU Hang detected"));

// 2. Abort NCCL collectives so that the remaining processes can exit.

nccl_api->AbortCollectives();

}

🌿 One-minute fallback timer:

// Immediately after issuing the first abort, arm a one-minute fallback timer.

auto fallback_guard = watchdog->Watch("Post-abort fallback", absl::Minutes(1), []() {

std::abort();

// or HangWatchdog::Abort()

});

✨ The GPU performs the computation, while the watchdog timer runs on the host CPU. The timer itself lives on the heap. Calling watchdog->Watch(...) registers a timer with the global HangWatchdog, which runs continuously on the CPU. When ExecuteThunksImpl starts, it asks the watchdog to create a ten-minute timer. The watchdog creates the timer and stores it in its internal queue. A background thread inside HangWatchdog continuously checks those timers. When the function returns, auto guard = watchdog->Watch(...) is destroyed. That stack variable is not the timer itself. It is only a handle to the timer stored on the heap. Destroying the guard signals the watchdog to cancel the corresponding timer.

📝 The GPU is a separate device, so constantly polling it from the CPU to check whether it has finished would be inefficient. Instead, ExecuteThunksImpl submits work through the execution stream. Along with the submitted work, it creates a stream event or future that will be signaled when the GPU has finished. A separate CPU synchronization thread waits for that notification. If the GPU stalls, the completion event is never signaled. The Hang Watchdog does not monitor the GPU directly. Instead, it monitors the CPU thread that is waiting for completion. If that thread remains blocked for too long, the watchdog invokes on_timeout().

🌱 Why not attach the watchdog to the waiting thread instead? That works for synchronous execution because the waiting happens inside ExecuteThunksImpl. In asynchronous execution, however, ExecuteThunksImpl never waits for completion. It only enqueues the work and immediately hands the completion future to RunAsync, which continues waiting on a different thread. The waiting thread no longer belongs to ExecuteThunksImpl.

🌼 My own analogy: the Hang Watchdog is like a police station, ExecuteThunksImpl is an employment agency, and the GPU is the worker. The agency receives jobs from either a synchronous company or an asynchronous company. The synchronous company asks the agency both to assign work and to keep monitoring the worker until the job is finished. The asynchronous company works differently. It only asks the agency to submit the job and immediately return. The problem is that once the agency leaves, it no longer has anyone monitoring the worker. If no one is responsible for keeping the watchdog alive during the asynchronous execution, the headquarters never receives the timeout signal either. The watchdog disappears simply because its lifetime ended with ExecuteThunksImpl, even though the GPU work was still running.

🍀 The watchdog guard is created inside ExecuteThunksImpl for both synchronous and asynchronous execution. The difference between the two modes is whether block_host_until_done is enabled. Simply forcing the guard to behave synchronously would not solve the problem because the guard itself is still a local stack variable. According to C++ object lifetime rules, it is destroyed as soon as ExecuteThunksImpl returns.