NVIDIA AI Engineer Interview Questions & Answers

11. Write a program to perform efficient batch tensor operations from scratch.CodingMediumNvidia

Question Details

In Python with NumPy, implement batched_matmul(left, right) for left shaped (batch_left, m, k) and right shaped (batch_right, k, n). The batch sizes must be equal, or one batch size may be one and broadcast across the other batch; no other broadcasting is allowed. Return shape (max(batch_left, batch_right), m, n) without a Python loop over batches and without materializing repeated copies of a broadcast operand. Require matching inner dimensions, finite floating values, the same float32 or float64 dtype and nonempty matrix dimensions; do not mutate inputs. State O(batch * m * k * n) arithmetic work and output-space requirements. Example: left=[[[1,2]]] and right=[[[3],[4]]] return [[[11]]].

Short Interview Answer (30-60 seconds)

I would first validate both 3D tensors. They must use the same float32 or float64 dtype, contain only finite values, and have matching inner dimensions. I allow only equal batch sizes or one batch size equal to 1. I then use np.broadcast_to to create views without repeated copies and perform one vectorized matrix multiplication. The arithmetic work is O(B * m * k * n). Auxiliary space is O(1) excluding the required O(B * m * n) output.

Detailed Explanation

See the Code while reading this explanation.

The function receives two batches of matrices. The left input has shape (batch_left, m, k). The right input has shape (batch_right, k, n). We multiply corresponding matrices and return a new batch with shape (B, m, n), where B = max(batch_left, batch_right). If one batch size is 1, that matrix can be reused across the other batch without making repeated copies. The inputs stay unchanged. NumPy can perform all batch multiplications in one vectorized operation.

Useful Questions to Ask the Interviewer

Should I reject every kind of broadcasting except equal batch sizes or one batch size equal to 1?
Should both inputs have exactly the same float32 or float64 dtype?
Should NaN and infinity values be rejected before multiplication?

Write a program to perform efficient batch tensor operations from scratch. diagram

How to Explain It in an Interview

1. Validate the input contract

Both inputs must be 3D NumPy arrays. The left shape is (batch_left, m, k). The right shape is (batch_right, k, n). Their dtypes must be equal and must be float32 or float64. Every value must be finite. The inner dimensions must match, so the left k must equal the right k. The matrix dimensions m, k, and n must be greater than zero. The function does not change either input.

2. Check the batch rule

Let bl be batch_left and br be batch_right. The batch sizes are valid when bl == br, bl == 1, or br == 1. No other broadcasting is allowed. Set B = max(bl, br). This is the batch size of the result.

3. Create broadcasted views

The left tensor is viewed with shape (B, m, k). The right tensor is viewed with shape (B, k, n). I use np.broadcast_to for this step. A broadcasted view lets NumPy reuse a batch-size-one matrix without materializing repeated copies.

4. Walk through the example

The example uses left = [[[1, 2]]] with shape (1, 1, 2). It uses right = [[[3], [4]]] with shape (1, 2, 1). The batch sizes already match, so B = 1. We multiply the only 1 by 2 row by the 2 by 1 column. The calculation is 1 * 3 + 2 * 4 = 11. The returned tensor has shape (1, 1, 1) and value [[[11]]].

5. Perform one vectorized multiplication

After creating the views, the code evaluates left_v @ right_v. NumPy treats the first dimension as the batch dimension and multiplies the last two dimensions as matrices. There is no Python loop over batches. The result has shape (B, m, n).

6. Explain why it is correct and state complexity

After broadcasting, every batch position has a left matrix with shape (m, k) and a right matrix with shape (k, n). Therefore each batch position produces a valid matrix with shape (m, n). If one input has batch size 1, its same matrix is reused at every required batch position. The arithmetic work is O(B * m * k * n). The returned array needs O(B * m * n) space. The broadcasted operands are views, so repeated copies are not allocated.

Key Insight / Why This Solution Works

The key idea is to separate validation, batch alignment, and multiplication. First, reject inputs that do not satisfy the exact contract. Next, set B = max(batch_left, batch_right). Use np.broadcast_to so both operands have batch size B without materializing repeated batch data. Finally, perform one vectorized matrix multiplication. The central invariant is that before multiplication the left view has shape (B, m, k) and the right view has shape (B, k, n). Therefore each corresponding pair of matrices has compatible inner dimensions and produces one (m, n) output matrix.

Code

from __future__ import annotations

from typing import Final

import numpy as np


def batched_matmul(left: np.ndarray, right: np.ndarray) -> np.ndarray:
    """Multiply batches of matrices with broadcasting only on the batch axis.

    left:  shape (batch_left, m, k)
    right: shape (batch_right, k, n)
    return: shape (max(batch_left, batch_right), m, n)
    """
    # Both inputs must have one batch axis and two matrix axes.
    if left.ndim != 3 or right.ndim != 3:
        raise ValueError("expected 3D arrays: (batch, m, k) and (batch, k, n)")

    # Require exactly the same floating dtype on both inputs.
    # The problem accepts only float32 and float64.
    if left.dtype != right.dtype or left.dtype not in (np.float32, np.float64):
        raise TypeError("dtype must match and be float32 or float64")

    # Reject NaN and infinity because every input value must be finite.
    if not np.isfinite(left).all() or not np.isfinite(right).all():
        raise ValueError("inputs must contain only finite values")

    # Read the batch size and matrix dimensions from both tensors.
    batch_left, m, left_k = left.shape
    batch_right, right_k, n = right.shape

    # Matrix multiplication requires the two inner dimensions to match.
    if left_k != right_k:
        raise ValueError(f"inner dimensions must match: {left_k} vs {right_k}")

    # Allow only equal batches or broadcasting from a batch size of one.
    # This rejects every other batch-broadcasting case.
    if batch_left != batch_right and batch_left != 1 and batch_right != 1:
        raise ValueError(f"incompatible batch sizes: {batch_left} vs {batch_right}")

    # The matrix dimensions must contain at least one element per axis.
    if m <= 0 or left_k <= 0 or right_k <= 0 or n <= 0:
        raise ValueError("matrix dimensions must be nonempty")

    # The larger allowed batch size becomes the output batch size.
    batch: Final[int] = max(batch_left, batch_right)

    # Create broadcasted views only. A batch-size-one operand is reused
    # logically without materializing repeated copies of its matrix data.
    left_v = np.broadcast_to(left, (batch, m, left_k))
    right_v = np.broadcast_to(right, (batch, right_k, n))

    # Perform every batch matrix product in one vectorized NumPy operation.
    # The returned array has shape (batch, m, n), and inputs are unchanged.
    result = left_v @ right_v

    return result


def main() -> None:
    # Run the exact example shown in the approved diagram.
    left = np.array([[[1, 2]]], dtype=np.float32)
    right = np.array([[[3], [4]]], dtype=np.float32)

    # 1 * 3 + 2 * 4 = 11, so the result is [[[11]]].
    result = batched_matmul(left, right)
    print(result)
    print("shape:", result.shape)


if __name__ == "__main__":
    main()

Time & Space Complexity

Let B = max(batch_left, batch_right). There are B output matrices. Each output matrix has m * n values. Computing one value requires k multiply-and-add work. So the total arithmetic work is O(B * m * k * n). The returned tensor contains B * m * n values, so output space is O(B * m * n). np.broadcast_to creates views instead of repeated tensor copies. Excluding the required output and constant-size view metadata, auxiliary space is O(1).

Where it is used

This pattern is useful when the same matrix operation must be applied to many items at once. It appears in machine learning tensor processing, batched linear algebra, feature transformations, and other numerical workloads. It is especially useful when one matrix must be reused across many batch entries because broadcasting avoids storing repeated copies.

Why Interviewers Ask This

The interviewer is checking whether you understand tensor shapes, matrix multiplication, controlled broadcasting, and vectorized NumPy code. They also want to see whether you validate an exact input contract instead of silently accepting broader NumPy behavior. The problem tests memory awareness because the broadcast operand must not be copied repeatedly. It also checks whether you can reason correctly about dtype rules, finite values, mutation, output shape, complexity, and edge cases.

Common interview mistakes

A common mistake is allowing NumPy's general broadcasting instead of enforcing the exact batch rule. Another mistake is using np.repeat or another operation that materializes repeated copies of a batch-size-one operand. Candidates may forget to require the same dtype, accept unsupported integer tensors, or fail to reject NaN and infinity values. Another error is forgetting to check the inner dimensions before multiplication. Finally, using a Python loop over batches breaks the main implementation requirement.

Interview tip

Explain the shape contract first. Then say that the key idea is to use broadcasted views followed by one vectorized matrix multiplication. Walk through 1 * 3 + 2 * 4 = 11. Finish by separating arithmetic work, O(B * m * k * n), from the required O(B * m * n) output space.

Interviewer may ask next

What changes when one input has batch size 1 and the other has batch size B?

The algorithm does not change. B is the larger batch size. np.broadcast_to creates a view of the batch-size-one operand with batch shape B, so the same matrix is reused at each batch position without materializing B copies. The other operand already has batch size B. The multiplication still uses one vectorized operation. Arithmetic work remains O(B * m * k * n), output space remains O(B * m * n), and auxiliary storage remains O(1) excluding the required output.

Why should I use np.broadcast_to instead of np.repeat for the batch-size-one operand?

np.repeat would materialize repeated tensor data. That violates the requirement and can add O(B * m * k) extra memory for the left operand or O(B * k * n) extra memory for the right operand. np.broadcast_to creates a view instead. The arithmetic work remains O(B * m * k * n), while the required result still uses O(B * m * n) output space.

12. What was your main goal when you were optimizing a ML model in this experience?BehavioralEasyNvidia

Question Details

Use a real project to state the user or system objective, baseline, candidate's ownership, optimization choices, guardrails, measured result, and what evidence changed the plan.

Interview tip:

Use STAR to structure your answer: briefly explain the Situation and Task, make Action the most detailed part, and finish with the Result. For example, describe a real model optimization project where you defined the main user or system goal, understood the baseline, owned the evaluation work, chose changes based on evidence, protected important guardrails, measured the result, and changed your plan when the data showed a different problem.

Situation

In my last role, I worked on a machine learning model that was giving useful results overall, but it was still making important mistakes on some real user cases. The main goal was not simply to make the model look better on a general test set. We wanted to improve the quality of the decisions that mattered to users while keeping the system reliable in production.

Task

I was responsible for understanding the current baseline, finding the main sources of error, and deciding which optimization work was most likely to improve the user outcome. I also needed to make sure that any change did not create new problems in areas such as reliability, response time, or important error cases.

Action

I first created a clear baseline using a held out evaluation set, which is data kept separate so we can test the model fairly. I grouped the model mistakes by type instead of looking only at one overall score. This showed me where the model was actually failing. My first idea was to spend most of the effort tuning the model settings. However, the error review showed that several important failures came from weak or inconsistent training examples. Based on that evidence, I changed the plan. I focused first on improving the relevant data and fixing unclear labels. I then tested model changes against the same baseline so the comparison stayed fair. For each candidate, I checked the main quality measure and also watched the important guardrails, including reliability and response time. I shared the findings with the team and explained why I was changing the original optimization plan. This kept the work focused on the user goal instead of optimizing a score that did not represent the real problem.

Result

The final approach improved the model on the important failure cases we had identified while keeping the production guardrails acceptable. The evaluation also gave us clearer evidence about which changes were useful and which were not. My main learning was that model optimization should start with the real user objective and careful error analysis. The best next step is not always a more complex model. Sometimes the strongest improvement comes from better data and a better way to measure the problem.

Why Interviewers Ask This

Interviewers ask this question to see whether the candidate can connect model optimization to a real user or system goal. A strong answer shows that the candidate can define a baseline, use evidence to choose the next step, protect important production constraints, change direction when the evidence disagrees with the original plan, and take clear ownership of the work.

Interviewer may ask next

Why did you change your plan from model tuning to improving the data?

I changed the plan because the error analysis showed that many important mistakes were connected to weak or inconsistent training examples. Tuning the model without fixing that problem would have treated the symptom instead of the cause. The evidence made data quality the higher value next step.

How did you decide whether the optimization was successful?

I compared each candidate with the same baseline on the held out evaluation set. I checked whether the important failure cases improved, not only whether one overall score moved. I also checked the production guardrails, such as reliability and response time. I considered the change successful only when it improved the target behavior without causing an unacceptable regression elsewhere.

13. Describe a challenging technical problem you solved in a previous project involving generative AI.BehavioralMediumNvidia

Question Details

Make the production symptom, constraints, candidate's ownership, hypotheses, architecture or model change, collaboration, validation, outcome, and lessons traceable.

Interview tip:

Use STAR to structure your answer: briefly explain the Situation and Task, make Action the most detailed part, and finish with the Result. For example, describe a production generative AI problem, the symptom you observed, the constraints you had, the part you owned, the possible causes you tested, the architecture or model change you made, how you worked with others, how you validated the fix, the outcome, and what you learned.

Situation

During a previous project, my team had a generative AI assistant that answered questions using internal documents. In testing, the answers looked good, but in production we started seeing an important reliability problem. Some answers sounded confident even when the retrieved documents did not contain enough evidence. We could not simply make the model refuse more often because users still needed useful answers for valid questions.

Task

I owned the investigation of the answer generation path. My goal was to find where unsupported answers were coming from and improve reliability without making the assistant much less useful. I also needed to work within the existing system instead of replacing the whole architecture.

Action

I first made the problem observable. I collected representative failure cases and traced each one through retrieval, prompt construction, and generation. This helped me separate two different causes. In some cases, retrieval returned weak or unrelated passages. In other cases, retrieval was reasonable, but the model still added details that were not supported by the context. I then tested these hypotheses one at a time. For the retrieval problem, I improved the filtering logic so low quality context was less likely to reach the model. For the generation problem, I changed the prompt so the model had a clear rule: answer from the provided evidence, and say that there is not enough information when the evidence is missing. I also added a validation step that checked whether important claims were supported by the retrieved context before the answer was returned. I worked with the application team to review examples where stricter validation could make useful answers disappear. We agreed on cases where the assistant should answer, ask for clarification, or state that it lacked enough evidence. I then compared the updated system with the earlier version using the same set of realistic questions. I reviewed both answer usefulness and whether claims stayed grounded in the source material. This showed that the combined retrieval, prompt, and validation changes addressed the real failure path better than changing only the model instructions.

Result

The assistant became more reliable because unsupported answers were caught earlier and weak evidence was handled more carefully. The team also had a clearer way to diagnose similar issues because we could see whether a failure started in retrieval, generation, or validation. I learned that a generative AI production problem should not automatically be treated as a model problem. I now trace the full path first, form separate hypotheses, and validate each change against realistic examples before changing the architecture.

Why Interviewers Ask This

Interviewers ask this question to see how you solve difficult AI engineering problems when the first cause is not obvious. A strong answer shows ownership, structured debugging, sound technical judgment, collaboration, careful validation, and the ability to learn from a production problem instead of only applying a quick model change.

Interviewer may ask next

Why did you change retrieval, prompting, and validation instead of only changing the prompt?

The failure did not have one cause. Some bad answers started with weak retrieved context, while others came from the model adding unsupported details even when the context was reasonable. A prompt change could help the second case, but it would not fix poor evidence entering the model. I changed each part only after tracing examples and confirming where the failure started.

What would you do differently if you faced the same problem again?

I would add the tracing and grounded answer checks earlier in the project. They made it much easier to tell whether a problem came from retrieval, prompt construction, generation, or validation. I would also build a small set of difficult production style examples before launch so the team could test these failure cases continuously instead of discovering them only after release.

14. Tell me about a time you had to troubleshoot a severe performance bottleneck in a distributed system under tight deadlines.BehavioralHardNvidia

Question Details

Use a production AI or GPU workload to show symptoms, containment, instrumentation, competing hypotheses, cross-team decisions, root cause, fix, verification, and prevention under the deadline.

Interview tip:

Use STAR to structure your answer: briefly explain the Situation and Task, make Action the most detailed part, and finish with the Result. For example, describe a production AI workload where you contained the immediate impact, added focused instrumentation, tested competing causes, coordinated with other teams, found the root cause, applied a safe fix, verified recovery, and added prevention steps before the deadline.

Situation

In my last role, I supported a distributed AI inference service that used several GPU workers. Shortly before an important delivery deadline, requests began taking much longer than normal. The GPUs were not consistently busy, while request queues kept growing. This told us that simply adding more GPU capacity might not solve the real problem.

Task

I was responsible for helping restore stable performance without creating a larger production risk. I needed to find where time was being lost, keep the service usable while we investigated, work with the platform team on shared infrastructure, and make a safe decision before the deadline.

Action

I first reduced the immediate pressure on the system. I worked with the team to limit nonessential traffic and protect the highest priority requests. This gave us room to investigate without allowing the queues to grow without control. I then traced one request through the main stages of the service. I compared time spent waiting in the request queue, preparing input data, sending work between services, running GPU inference, and returning the result. I also checked GPU utilization, worker queue depth, CPU use, memory pressure, and network behavior. I kept several explanations open instead of assuming the model was slow. The main possibilities were overloaded GPUs, slow input preparation, network delays, or workers waiting on shared resources. The measurements showed an important pattern. GPU execution itself was healthy when work reached the device, but workers were spending too much time waiting before inference started. I narrowed the problem to a shared data preparation path that had become a bottleneck as traffic increased. I reproduced the behavior under controlled load and shared the evidence with the platform team. We discussed two choices. One was a larger infrastructure change that could improve the design but carried more risk under the deadline. The other was a smaller change that reduced contention in the shared path and could be tested quickly. I recommended the smaller fix first because it addressed the measured bottleneck and was easier to reverse if something went wrong. We tested it under representative load, compared queue behavior and GPU activity before and after the change, and then rolled it out carefully. I watched the same signals during the rollout so we could stop if performance became worse.

Result

The service returned to stable behavior, request queues stopped building up, and the GPUs received work more consistently. We met the delivery need without making a risky redesign during the incident. Afterward, I helped add better stage level timing, queue monitoring, and load tests so the same kind of bottleneck would be easier to detect earlier. The main lesson for me was to measure each stage of a distributed system before adding capacity. A low GPU utilization number does not always mean the GPU is the problem.

Why Interviewers Ask This

Interviewers ask this question to see how a candidate behaves during a high pressure production problem. They want evidence that the candidate can contain impact, use measurements instead of guesses, compare several possible causes, make safe tradeoffs under time pressure, work across teams, verify a fix, and improve the system after the incident.

Interviewer may ask next

Why did you choose the smaller fix instead of making the larger infrastructure change immediately?

The measurements showed that the smaller fix directly addressed the current bottleneck. It was also easier to test and reverse. Because we were under a tight deadline, I wanted to reduce production risk while solving the measured problem. I supported considering the larger design change later, when we had time to test it properly.

What would you do differently if you faced a similar performance incident now?

I would make stage level timing and queue monitoring part of the service from the beginning. In this incident, we had to add some of that visibility while troubleshooting. Having it ready earlier would let me separate GPU execution time from waiting time much faster. I would also run regular load tests that stress shared paths, because a component that works well at normal traffic can become the limiting point as concurrency grows.

Questions with Detailed ExplanationsWith Detailed Explanations

11. Write a program to perform efficient batch tensor operations from scratch.CodingMediumNvidia

12. What was your main goal when you were optimizing a ML model in this experience?BehavioralEasyNvidia

13. Describe a challenging technical problem you solved in a previous project involving generative AI.BehavioralMediumNvidia

14. Tell me about a time you had to troubleshoot a severe performance bottleneck in a distributed system under tight deadlines.BehavioralHardNvidia